Skip problematic batches if grad nan is detected.
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
During SFT training, we sometimes encounter errors caused by gradient NaNs. After re-running, the issue is usually resolved, so I think it’s better to skip the problematic batches whenever a gradient NaN is detected.
```
File "/opt/NeMo-RL/nemo_rl/utils/nsys.py", line 88, in wrapper
ret = func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/NeMo-RL/nemo_rl/models/policy/megatron_policy_worker.py", line 902, in train
losses_reduced = forward_backward_func(
^^^^^^^^^^^^^^^^^^^^^^
File "/opt/NeMo-RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 543, in forward_backward_no_pipelining
backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
File "/opt/NeMo-RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 414, in backward_step
custom_backward(output_tensor[0], output_tensor_grad[0])
File "/opt/NeMo-RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 158, in custom_backward
Variable._execution_engine.run_backward(
File "/opt/NeMo-RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py", line 510, in hook
self.param_to_bucket_group[param].register_grad_ready(param)
File "/opt/NeMo-RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/param_and_grad_buffer.py", line 436, in register_grad_ready
self.start_grad_sync()
File "/opt/NeMo-RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/param_and_grad_buffer.py", line 293, in start_grad_sync
self.check_grads(
File "/opt/NeMo-RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/param_and_grad_buffer.py", line 165, in check_grads
rerun_state_machine.validate_result(
File "/opt/NeMo-RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/rerun_state_machine.py", line 505, in validate_result
raise RuntimeError(full_message)
RuntimeError: Rank 60, xxxx, device xxx, iteration 11: Unexpected result nan (message='found NaN in local grad norm for bucket #0 in backward pass before data-parallel communication collective')
```
**Steps/Code to reproduce bug**
Please list *minimal* steps or code snippet for us to be able to reproduce the bug.
A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.
**Expected behavior**
A clear and concise description of what you expected to happen.
**Additional context**
Add any other context about the problem here.
Contributor guide
Assessment
This issue has not been assessed yet.