NVIDIA-NeMo / NVIDIA-NeMo/RL

On policy distillation initializes optimizer states for teacher

Open
#2,737 0 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
bug community-request Memory waiting-on-maintainers
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

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.