Training stops after long checkpointing

**Describe the bug**

I am training a Qwen3-32B model with Gym. The training stops after the checkpointing operation (that took quite a long time on my machine, around 6 minutes). The issue did not happen when training a Qwen3-8B model.

**Steps/Code to reproduce bug**

A way to reproduce the bug is to use [this fork](https://github.com/owkin/Owkin-NeMo-RL) of Nemo-RL using also a [fork of NeMo-Gym](https://github.com/owkin/Owkin-Gym/tree/9adfa4d029f59e0c4f7ccd97d1d4c36b68ba4efa) because I am using custom resources servers.

I am using this [grpo config file](https://github.com/owkin/Owkin-NeMo-RL/blob/main/examples/nemo_gym/grpo_config_random.yaml) to run a grpo training on a single node with 8 H200 gpus.

Here is the error that I get (after step 2 when checkpointing occurs)

```
▶ Training policy...
/fsx/xav/code/OwkinZero/nemo-rl/examples/nemo_gym/run_grpo_nemo_gym.py:280: UserWarning: You asked to save checkpoints based on policy_idea_generator_agent/reward/mean but no val metrics were collected. This checkpoint will not be saved as top-k.
grpo_train(
Saving checkpoint for step 2...
(MegatronPolicyWorker[rank=0] pid=1750241) WARNING:megatron.core.utils:utils.py:890: make_tp_sharded_tensor_for_checkpoint received extra kwargs: ['allow_shape_mismatch']
(MegatronPolicyWorker[rank=0] pid=1750241) saving checkpoint at iteration 0 to /fsx/xav/code/OwkinZero/nemo-rl/results/grpo-llm-judge-qwen3-32b/tmp_step_2/policy/weights in torch_dist format
(MegatronPolicyWorker[rank=0] pid=1750241) Storing distributed optimizer sharded state of type dp_reshardable
(VllmAsyncGenerationWorker pid=1743926) (EngineCore_DP0 pid=1745720) (RayWorkerWrapper pid=1745909) INFO 01-23 10:15:26 [cumem.py:239] CuMemAllocator: sleep freed 110.62 GiB memory in total, of which 15.49 GiB is backed up in CPU and the rest 95.13 GiB is discarded directly. [repeated 4x across cluster]
(VllmAsyncGenerationWorker pid=1743926) (EngineCore_DP0 pid=1745720) (RayWorkerWrapper pid=1745907) INFO 01-23 10:11:02 [gpu_worker.py:142] Sleep mode freed 111.04 GiB memory, 6.57 GiB memory is still in use. [repeated 4x across cluster]
(VllmAsyncGenerationWorker pid=1743926) (EngineCore_DP0 pid=1745720) INFO 01-23 10:15:29 [abstract.py:306] It took 9.306633 seconds to fall asleep.
(MegatronPolicyWorker[rank=0] pid=1750241) successfully saved checkpoint from iteration 0 to /fsx/xav/code/OwkinZero/nemo-rl/results/grpo-llm-judge-qwen3-32b/tmp_step_2/policy/weights [ t 1/8, p 1/1 ]
(MegatronPolicyWorker[rank=0] pid=1750241) Saved checkpoint to /fsx/xav/code/OwkinZero/nemo-rl/results/grpo-llm-judge-qwen3-32b/tmp_step_2/policy/weights
Logged data to logs/grpo-llm-judge-qwen3-32b/exp_047/train_data_step2.jsonl

📊 Training Results:
• Loss: 0.0394
• Generation KL Error: 0.0008
• Avg Reward: 0.2993
• Mean Generation Length: 1726.7266

⏱️ Timing:
• Total step time: 653.13s
• checkpointing: 366.80s (56.2%)
• generation: 124.74s (19.1%)
• policy_training: 57.95s (8.9%)
• policy_and_reference_logprobs: 39.62s (6.1%)
• prepare_for_generation/total: 33.78s (5.2%)
• training_prep: 28.99s (4.4%)
• prepare_for_generation/transfer_and_update_weights: 1.30s (0.2%)
• logprob_inference_prep: 1.20s (0.2%)
• data_processing: 0.03s (0.0%)
• reward_calculation: 0.01s (0.0%)

🔍 Performance Metrics:
• Throughputs (per GPU):
- E2E (Samples/sec/gpu): 0.02
- E2E (Tokens/sec/gpu): 60.88
- Policy Training (Tokens/sec/gpu): 686.26
- Policy and Reference Logprobs (Tokens/sec/gpu): 1003.75
- Training Worker Group (Tokens/sec/gpu): 407.59
- Generation Worker Group (Tokens/sec/gpu): 318.78
• Throughputs (per Group):
- E2E (Samples/sec): 0.20
- E2E (Tokens/sec): 487.08
- Training Worker Group (Tokens/sec): 3260.74
- Generation Worker Group (Tokens/sec): 2550.24
• Training FLOPS: 1007.55 TFLOPS (125.94 TFLOPS per rank)
• Training Model Floating Point Utilization: 12.73%

========================= Step 3/101 =========================
▶ Preparing batch...
▶ Generating responses for batch of size 128...
(MegatronPolicyWorker[rank=0] pid=1750241) GPU Memory before optimizer offload: 53.80GB allocated, 54.20GB reserved
Traceback (most recent call last):
File "/fsx/xav/code/OwkinZero/nemo-rl/examples/nemo_gym/run_grpo_nemo_gym.py", line 297, in
main()
File "/fsx/xav/code/OwkinZero/nemo-rl/examples/nemo_gym/run_grpo_nemo_gym.py", line 280, in main
grpo_train(
File "/fsx/xav/code/OwkinZero/nemo-rl/nemo_rl/algorithms/grpo.py", line 1136, in grpo_train
refit_policy_generation(
File "/fsx/xav/code/OwkinZero/nemo-rl/nemo_rl/algorithms/grpo.py", line 924, in refit_policy_generation
policy.offload_before_refit()
File "/fsx/xav/code/OwkinZero/nemo-rl/nemo_rl/models/policy/lm_policy.py", line 785, in offload_before_refit
ray.get(futures)
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 2882, in get
values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 970, in get_objects
raise value
ray.exceptions.ActorUnavailableError: The actor 899dad2d0066f26d5ea5e02f01000000 is unavailable: The actor is temporarily unavailable: RpcError: RPC Error message: Socket closed; RPC Error details: rpc_code: 14. The task may or may not have been executed on the actor.
Traceback (most recent call last):
File "/fsx/xav/code/OwkinZero/nemo-rl/examples/nemo_gym/run_grpo_nemo_gym.py", line 297, in
main()
File "/fsx/xav/code/OwkinZero/nemo-rl/examples/nemo_gym/run_grpo_nemo_gym.py", line 280, in main
grpo_train(
File "/fsx/xav/code/OwkinZero/nemo-rl/nemo_rl/algorithms/grpo.py", line 1136, in grpo_train
refit_policy_generation(
File "/fsx/xav/code/OwkinZero/nemo-rl/nemo_rl/algorithms/grpo.py", line 924, in refit_policy_generation
policy.offload_before_refit()
File "/fsx/xav/code/OwkinZero/nemo-rl/nemo_rl/models/policy/lm_policy.py", line 785, in offload_before_refit
ray.get(futures)
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 2882, in get
values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 970, in get_objects
raise value
ray.exceptions.ActorUnavailableError: The actor 899dad2d0066f26d5ea5e02f01000000 is unavailable: The actor is temporarily unavailable: RpcError: RPC Error message: Socket closed; RPC Error details: rpc_code: 14. The task may or may not have been executed on the actor.
```

**Expected behavior**

Training should continue after checkpointing as it's the case when I'm using a 8B model:

```
▶ Training policy...
/fsx/xav/code/OwkinZero/nemo-rl/examples/nemo_gym/run_grpo_nemo_gym.py:280: UserWarning: You asked to save checkpoints based on policy_idea_generator_agent/reward/mean but no val metrics were collected. This checkpoint will not be saved as top-k.
grpo_train(
Saving checkpoint for step 2...
(MegatronPolicyWorker[rank=0] pid=3768308) WARNING:megatron.core.utils:utils.py:890: make_tp_sharded_tensor_for_checkpoint received extra kwargs: ['allow_shape_mismatch']
(MegatronPolicyWorker[rank=0] pid=3768308) saving checkpoint at iteration 0 to /fsx/xav/code/OwkinZero/nemo-rl/results/grpo-llm-judge-qwen3-32b/tmp_step_2/policy/weights in torch_dist format
(MegatronPolicyWorker[rank=0] pid=3768308) Storing distributed optimizer sharded state of type dp_reshardable
(VllmAsyncGenerationWorker pid=3762026) (EngineCore_DP0 pid=3763826) INFO 01-23 11:27:05 [block_pool.py:390] Successfully reset prefix cache [repeated 3x across cluster]
(VllmAsyncGenerationWorker pid=3762026) (EngineCore_DP0 pid=3763826) (RayWorkerWrapper pid=3764022) INFO 01-23 11:27:07 [cumem.py:239] CuMemAllocator: sleep freed 110.51 GiB memory in total, of which 3.84 GiB is backed up in CPU and the rest 106.67 GiB is discarded directly. [repeated 4x across cluster]
(VllmAsyncGenerationWorker pid=3762026) (EngineCore_DP0 pid=3763826) (RayWorkerWrapper pid=3764023) INFO 01-23 11:23:52 [gpu_worker.py:142] Sleep mode freed 110.74 GiB memory, 6.32 GiB memory is still in use. [repeated 4x across cluster]
(VllmAsyncGenerationWorker pid=3762026) (EngineCore_DP0 pid=3763826) INFO 01-23 11:27:08 [abstract.py:306] It took 2.762003 seconds to fall asleep.
(MegatronPolicyWorker[rank=0] pid=3768308) successfully saved checkpoint from iteration 0 to /fsx/xav/code/OwkinZero/nemo-rl/results/grpo-llm-judge-qwen3-32b/tmp_step_2/policy/weights [ t 1/8, p 1/1 ]
(MegatronPolicyWorker[rank=0] pid=3768308) Saved checkpoint to /fsx/xav/code/OwkinZero/nemo-rl/results/grpo-llm-judge-qwen3-32b/tmp_step_2/policy/weights
Logged data to logs/grpo-llm-judge-qwen3-32b/exp_052/train_data_step2.jsonl

📊 Training Results:
• Loss: 0.1727
• Generation KL Error: 0.0008
• Avg Reward: 0.3361
• Mean Generation Length: 2232.3828

⏱️ Timing:
• Total step time: 295.93s
• generation: 125.02s (42.2%)
• checkpointing: 97.86s (33.1%)
• policy_training: 32.79s (11.1%)
• policy_and_reference_logprobs: 22.21s (7.5%)
• prepare_for_generation/total: 8.82s (3.0%)
• training_prep: 8.55s (2.9%)
• prepare_for_generation/transfer_and_update_weights: 0.69s (0.2%)
• logprob_inference_prep: 0.64s (0.2%)
• data_processing: 0.03s (0.0%)
• reward_calculation: 0.00s (0.0%)

🔍 Performance Metrics:
• Throughputs (per GPU):
- E2E (Samples/sec/gpu): 0.05
- E2E (Tokens/sec/gpu): 161.71
- Policy Training (Tokens/sec/gpu): 1459.56
- Policy and Reference Logprobs (Tokens/sec/gpu): 2154.57
- Training Worker Group (Tokens/sec/gpu): 870.12
- Generation Worker Group (Tokens/sec/gpu): 382.78
• Throughputs (per Group):
- E2E (Samples/sec): 0.43
- E2E (Tokens/sec): 1293.70
- Training Worker Group (Tokens/sec): 6960.96
- Generation Worker Group (Tokens/sec): 3062.26
• Training FLOPS: 563.35 TFLOPS (70.42 TFLOPS per rank)
• Training Model Floating Point Utilization: 7.12%

========================= Step 3/101 =========================
▶ Preparing batch...
▶ Generating responses for batch of size 128...
(MegatronPolicyWorker[rank=0] pid=3768308) GPU Memory before optimizer offload: 13.45GB allocated, 13.54GB reserved
(MegatronPolicyWorker[rank=3] pid=3768360) GPU Memory after optimizer offload: 2.00GB allocated, 2.10GB reserved
(MegatronPolicyWorker[rank=7] pid=3768736) Saved checkpoint to /fsx/xav/code/OwkinZero/nemo-rl/results/grpo-llm-judge-qwen3-32b/tmp_step_2/policy/weights [repeated 7x across cluster]
(MegatronPolicyWorker[rank=7] pid=3768736) GPU Memory before optimizer offload: 13.45GB allocated, 13.54GB reserved [repeated 7x across cluster]
(VllmAsyncGenerationWorker pid=3762051) (EngineCore_DP0 pid=3763830) INFO 01-23 11:29:58 [abstract.py:324] It took 0.329717 seconds to wake up tags ['weights'].
(MegatronPolicyWorker[rank=0] pid=3768308) MegatronPolicyWorker[rank=0]: Packed 1 groups of tensors
(MegatronPolicyWorker[rank=2] pid=3768335) GPU Memory after refit complete: 0.10GB allocated, 0.19GB reserved
(VllmAsyncGenerationWorker pid=3762051) (EngineCore_DP0 pid=3763830) INFO 01-23 11:30:00 [abstract.py:324] It took 0.200763 seconds to wake up tags ['kv_cache'].
Collecting rollouts: 0%| | 0/128 [00:00