CPU offload issue causing optimizer to work on randomly initialized model weights
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
When using megatron backend + optimizer CPU offload + fp32 optimizer update + distributed optimizer + reading from a hf checkpoints combined together, the model weights that the optimizer states operate on (call step() on) are still randomly initiated values invalidating the training.
**Steps/Code to reproduce bug**
In this [PR](https://github.com/NVIDIA-NeMo/RL/pull/2372) which is aimed to fix the issue has steps to reproduce and fix the issue.
run:
`NRL_DISABLE_FP32_SYNC_FIX=1 bash tests/test_suites/llm/sft-qwen2.5-0.5b-1n2g-megatron-cpuoffload.sh`
loss values ~10 which is logprobs of random prediction of ~100K vocab
When calling with
`NRL_DISABLE_FP32_SYNC_FIX=0 bash tests/test_suites/llm/sft-qwen2.5-0.5b-1n2g-megatron-cpuoffload.sh`
**Root Cause**
The root cause is after when the above mentioned knobs are combined, some more copies on CPU of the model weights are created, however, they are not correctly replaced by the pretrained hf model weights when doing fine-tuning, causing the optimizer to operate on randomly initialized model weights invalidating the whole training.
**Expected behavior**
A clear and concise description of what you expected to happen.
**Additional context**
The issue was first discovered during a mutli-node large model fine tuning where the fine-tuned model only produced gibberish outputs, further investigation revealed that the model weights were nearly randomly initialized, then this final root cause was discovered. The cpu offload + distributed optimizer + bf16 are all combined for efficiency.
Contributor guide
Assessment
This issue has not been assessed yet.