perf: reduce .item() calls and fuse optimizer - #150
Conversation
Signed-off-by: Doğaç Eldenk <dogacel@gmail.com>
Signed-off-by: Doğaç Eldenk <dogacel@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 665b03169c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| grad_norm = torch.nn.utils.clip_grad_norm_(self.fp32_params, self.max_grad_norm) | ||
| if grad_norm > 0.0: | ||
| self.optimizer.step() | ||
| self.optimizer.step() |
There was a problem hiding this comment.
When a training batch produces NaN gradients, clip_grad_norm_ returns NaN; the removed grad_norm > 0.0 guard previously skipped the AdamW update because the comparison was false. With the unconditional step here, AdamW consumes the NaN gradients and the fp32 master parameters are then copied back into the bf16 model, so a single bad batch can poison subsequent training and checkpoints instead of being skipped.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
echo one this, let's keep the guard
There was a problem hiding this comment.
sorry I missed this
@Dogacel may you fix in the followup
3 participants
Improves training speed about ~25%.