gemma-3-4b-it has abnormal behavior with load_checkpoint
- 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
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