On policy distillation initializes optimizer states for teacher
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
On policy distillation initializes optimizer states for teacher
**Steps/Code to reproduce bug**
Running OPD with megatron following the recipe, for example `examples/configs/distillation_math_megatron.yaml`
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**
Teacher doesn't initialize optimizer states, so that GPU VRAM is not filled with extra memory.
**Additional context**
I have investigated the issue and the reason is because,
in file `nemo_rl/models/megatron/setup.py`, get_model method is initialized without setting any `wrap_with_ddp`, but by default it is set to True by megatron bridge, see file `Megatron-Bridge/src/megatron/bridge/models/model_provider.py`.
The fix is to set `wrap_with_ddp=load_optimizer` in file `nemo_rl/models/megatron/setup.py`:
```python
model = get_model(
megatron_cfg.model,
megatron_cfg.ddp,
use_torch_fsdp2=megatron_cfg.dist.use_torch_fsdp2,
overlap_param_gather_with_optimizer_step=megatron_cfg.optimizer.overlap_param_gather_with_optimizer_step,
data_parallel_random_init=megatron_cfg.rng.data_parallel_random_init,
pre_wrap_hook=pre_wrap_hook,
mixed_precision_wrapper=mixed_precision_wrapper,
# patch: to avoid creating optimizer for teacher model
wrap_with_ddp=load_optimizer,
pg_collection=pg_collection,
)
```
Contributor guide
Assessment
This issue has not been assessed yet.