NVIDIA-NeMo / NVIDIA-NeMo/RL

gemma-3-4b-it has abnormal behavior with load_checkpoint

Open
#907 0 comments 0 reactions 0 assignees View on GitHub
bug t-pytdensor
Dominant language
Python
Stars
2k
Forks
562
Avg merge
4d 5h
Merged PRs (30d)
145

Description

Thanks @RayenTian for pointing out this when testing `gemma-3-4b-it`.
It happens on `gemma-3-4b-it` (or maybe other VLMs), pure LMs like `Qwen2.5-1.5B` haven't this issue.

Although the error log shows it's missing some refit params, it is because the `__init__` function is somehow skip or hang at `self.load_checkpoint(weights_path, optimizer_path)`. You can print something after this line to confirm.

Move these refit params upper than `self.load_checkpoint(weights_path, optimizer_path)` can run through, but I think the optimizer isn't successfully inited in this way.
https://github.com/NVIDIA-NeMo/RL/blob/e1f56c42ae175d3863ccaf4e21b7de7e9c46c2e1/nemo_rl/models/policy/dtensor_policy_worker.py#L416-L431

Have some investigate at `load_checkpoint`.
1. `dcp.load(state_dict=model_state_dict, checkpoint_id=weights_path)` works fine.
2. `dcp.load(state_dict=optimizer_state_dict, checkpoint_id=optimizer_path)` is skip or hang. You can print something after this line to confirm.
https://github.com/NVIDIA-NeMo/RL/blob/e1f56c42ae175d3863ccaf4e21b7de7e9c46c2e1/nemo_rl/utils/native_checkpoint.py#L191-L202

**Repro**
1. Run the script first time to get a checkpoint.
2. Run the script again to reproduce.
```bash
uv run python examples/run_grpo_math.py \
grpo.max_num_steps=2 \
policy.model_name='google/gemma-3-4b-it' \
policy.dynamic_batching.enabled=False \
policy.sequence_packing.enabled=True \
policy.max_total_sequence_length=2048 \
policy.dtensor_cfg.tensor_parallel_size=4 \
policy.dtensor_cfg.activation_checkpointing=True \
policy.train_micro_batch_size=1 \
checkpointing.checkpoint_dir='results/gemma-3-4b-it-tp4' \
checkpointing.save_period=100 \
checkpointing.enabled=true \
logger.wandb_enabled=false \
logger.tensorboard_enabled=false \
cluster.gpus_per_node=4 \
cluster.num_nodes=1
```

**Error log**
```bash
Traceback (most recent call last):
File "/home/scratch.yukih_gpu/depot/reinforcer/examples/run_grpo_math.py", line 272, in
main()
File "/home/scratch.yukih_gpu/depot/reinforcer/examples/run_grpo_math.py", line 255, in main
grpo_train(
File "/home/scratch.yukih_gpu/depot/reinforcer/nemo_rl/algorithms/grpo.py", line 566, in grpo_train
refit_policy_generation(
File "/home/scratch.yukih_gpu/depot/reinforcer/nemo_rl/algorithms/grpo.py", line 444, in refit_policy_generation
ipc_handles = policy.get_weights_ipc_handles(keys)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/scratch.yukih_gpu/depot/reinforcer/nemo_rl/models/policy/lm_policy.py", line 548, in get_weights_ipc_handles
worker_handles: list[dict[str, Any]] = ray.get(
^^^^^^^^
File "/home/scratch.yukih_gpu/depot/reinforcer/.venv/lib/python3.12/site-packages/ray/_private/auto_init_hook.py", line 21, in auto_init_wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/home/scratch.yukih_gpu/depot/reinforcer/.venv/lib/python3.12/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/home/scratch.yukih_gpu/depot/reinforcer/.venv/lib/python3.12/site-packages/ray/_private/worker.py", line 2822, in get
values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/scratch.yukih_gpu/depot/reinforcer/.venv/lib/python3.12/site-packages/ray/_private/worker.py", line 930, in get_objects
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AttributeError): ray::DTensorPolicyWorker.get_weights_ipc_handles() (pid=298727, ip=192.168.1.1, actor_id=7a62ef9fe916defa4ad3691101000000, repr=DTensorPolicyWorker[rank=2])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/scratch.yukih_gpu/depot/reinforcer/venvs/nemo_rl.models.policy.dtensor_policy_worker.DTensorPolicyWorker/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/home/scratch.yukih_gpu/depot/reinforcer/nemo_rl/utils/nsys.py", line 88, in wrapper
ret = func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/home/scratch.yukih_gpu/depot/reinforcer/nemo_rl/models/policy/dtensor_policy_worker.py", line 1249, in get_weights_ipc_handles
if self._held_streamed_param_reference is not None:
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'DTensorPolicyWorker' object has no attribute '_held_streamed_param_reference'
```

Contributor guide

Open the contributing guide

Research direction

Start with nemo_rl/models/policy/dtensor_policy_worker.py around the referenced __init__ and load_checkpoint call, then inspect nemo_rl/utils/native_checkpoint.py around the optimizer dcp.load. Run the provided gemma-3-4b-it command once to create a checkpoint and again to reproduce the hang or skip. Done means the checkpoint reload completes with optimizer state initialized and the subsequent get_weights_ipc_handles call no longer raises AttributeError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.