Skip to content

GitLab

Commit 4fd6432d authored by Lawrence McAfee's avatar Lawrence McAfee
Browse files

fixed zero_grad for fp32_from_float16_groups

parent 90e0a0dd
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
megatron/optimizer/optimizer.py
View file @ 4fd6432d
...@@ -285,6 +285,8 @@ class Float16OptimizerWithFloat16Params(MegatronOptimizer): ...@@ -285,6 +285,8 @@ class Float16OptimizerWithFloat16Params(MegatronOptimizer):
float16_groups & fp32_from_fp32_groups.""" float16_groups & fp32_from_fp32_groups."""
for group in self.float16_groups: for group in self.float16_groups:
_zero_grad_group_helper(group, set_to_none) _zero_grad_group_helper(group, set_to_none)
for group in self.fp32_from_float16_groups:
_zero_grad_group_helper(group, set_to_none)
for group in self.fp32_from_fp32_groups: for group in self.fp32_from_fp32_groups:
_zero_grad_group_helper(group, set_to_none) _zero_grad_group_helper(group, set_to_none)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment