NVIDIA-NeMo / NVIDIA-NeMo/RL

Skip problematic batches if grad nan is detected.

Open
#1,230 1 comment 0 reactions 1 assignee Claimed by @yaoyu-33 View on GitHub
bug research
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.