feat: Gradient clipping and accumulation - #660
Conversation
|
|
ceaf307 to
f1d1c58
Compare
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #660 +/- ##
==========================================
+ Coverage 85.58% 85.72% +0.14%
==========================================
Files 258 258
Lines 11378 11392 +14
==========================================
+ Hits 9738 9766 +28
+ Misses 1640 1626 -14 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
MaxiBoether
left a comment
MaxiBoether
left a comment
There was a problem hiding this comment.
Thanks John, a couple of small comments
| @field_validator("gradient_accumulation_steps") | ||
| @classmethod | ||
| def validate_accumulation_steps(cls, value: int) -> int: | ||
| if value < 1: | ||
| raise ValueError("Invalid value for gradient_accumulation_steps, must be greater than 0.") | ||
| return value | ||
|
|
||
| @field_validator("max_grad_norm") | ||
| @classmethod | ||
| def validate_max_grad_norm(cls, value: float) -> float: | ||
| if value is not None and value <= 0: | ||
| raise ValueError("Invalid value for max_grad_norm, must be greater than 0.") |
There was a problem hiding this comment.
this should be checked via the min/max value of Field, not via extra validators
| ) | ||
| max_grad_norm: float | None = Field( | ||
| default=None, | ||
| description="Clips the gradients normed over this value, if its 0 it will not be used.", |
There was a problem hiding this comment.
You mean if it's None nothing will be clipped?
There was a problem hiding this comment.
If it is None then nothing is clipped, we skip the function in the pytorch trainer, I guess we can also just skip when it is 0 since it never makes sense to clip to 0.
| bool enable_accurate_gpu_measurements = 25; | ||
| int64 record_loss_every = 26; | ||
| bool drop_last_batch = 27; | ||
| float grad_norm = 28; |
There was a problem hiding this comment.
can we rename this to max_grad_norm?
There was a problem hiding this comment.
and shouldn't it be optional float? since it's optional in the config schema?
| # Instead of checking the logged "num_batches_trained" (which counts mini-batches), | ||
| # we compare the final model parameters to ensure they are equivalent. | ||
| final_params_accum = clone_params(trainer_accum._model.model) | ||
| final_params_single = clone_params(trainer_single._model.model) | ||
|
|
||
| for name in init_params.keys(): | ||
| assert torch.allclose(final_params_accum[name], final_params_single[name], atol=1e-5), f"Mismatch in {name}" | ||
|
|
There was a problem hiding this comment.
I appreciate this test, but can we somehow check the actual accumulation and not the side effect (same parameters)? e.g., can we mock the optimizer and model to check how often this/with which data they were called? we can keep this test and add another test, but right now we don't actually test the accumulation
There was a problem hiding this comment.
Made a new test, I think it does what you wanted.
| self._max_grad_norm = training_info.max_grad_norm | ||
| self.gradient_accumulation_steps = training_info.gradient_accumulation_steps |
There was a problem hiding this comment.
I think if you make max grad norm private then also gradient accumulation steps should be private. Can you move those to lines to line ~ 121 (where we set a lot of internal variables)?
MaxiBoether
changed the title
fe01aa7 to
efdf0b9
Compare
2 participants
Added gradient clipping and accumulation.