CUDA OOM during initial HF→Megatron checkpoint conversion in MegatronPolicyWorker initialization

**Describe the bug**

When using the Megatron backend in NeMo-RL, the policy worker fails during initialization when importing a Hugging Face model for the first time. Additionally, a similar failure can also be reproduced with the DTensor backend, suggesting that the issue may not be specific to Megatron checkpoint conversion alone.

The failure occurs during the HF → Megatron checkpoint conversion process. Specifically, the crash happens inside Megatron-Bridge checkpoint saving, where tensors are being transferred from GPU to CPU during async checkpoint preparation.

As a result, the MegatronPolicyWorker actor dies during creation and training never starts.
**Steps/Code to reproduce bug**

Please list *minimal* steps or code snippet for us to be able to reproduce the bug.

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**

Use a Windows host with WSL2.
Pull the official NeMo-RL container:
docker pull nvcr.io/nvidia/nemo-rl:v0.6.0

Start the container with GPU access.
Enter the container and launch training with:

uv run examples/run_grpo.py \
--config examples/configs/grpo_math_1B_megatron.yaml

The failure occurs during the initial Hugging Face → Megatron checkpoint conversion stage, before training begins.

**Additional context**

root@7051ad57b771:/opt/nemo-rl# uv run examples/run_grpo.py --config examples/configs/grpo_math_1B_megatron.yaml
/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/fsdp/src/megatron_fsdp/mixed_precision.py:115: UserWarning: Transformer Engine and Apex are not installed. Falling back to local implementations of multi_tensor_applier and multi_tensor_scale
warnings.warn(
/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/optimizer/__init__.py:25: UserWarning: Transformer Engine and Apex are not installed. Falling back to Torch optimizers.
warnings.warn(
/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/optimizer/optimizer.py:29: UserWarning: Transformer Engine and Apex are not installed. Falling back to local implementations of multi_tensor_applier and multi_tensor_scale
warnings.warn(
/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/optimizer/clip_grads.py:32: UserWarning: Transformer Engine and Apex are not installed. Falling back to local implementations of multi_tensor_applier, multi_tensor_l2norm, and multi_tensor_scale
warnings.warn(
/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/models/backends.py:34: UserWarning: Apex is not installed. Falling back to TorchNorm
warnings.warn("Apex is not installed. Falling back to Torch Norm")
Loaded configuration from: examples/configs/grpo_math_1B_megatron.yaml
Applied CLI overrides
Final config:
{'checkpointing': {'checkpoint_dir': 'results/grpo_megatron',
'checkpoint_must_save_by': None,
'enabled': False,
'higher_is_better': True,
'keep_top_k': 3,
'metric_name': 'val:accuracy',
'model_save_format': 'safetensors',
'save_consolidated': False,
'save_optimizer': True,
'save_period': 10},
'cluster': {'gpus_per_node': 1, 'num_nodes': 1},
'data': {'default': {'env_name': 'math',
'processor': 'math_hf_data_processor',
'prompt_file': 'examples/prompts/cot.txt',
'system_prompt_file': None},
'max_input_seq_length': 512,
'num_workers': 1,
'shuffle': True,
'train': {'dataset_name': 'OpenMathInstruct-2',
'seed': 42,
'split_validation_size': 0.05},
'use_multiple_dataloader': False,
'validation': None},
'env': {'math': {'math_verify_impl': 'hf_math_verify', 'num_workers': 8}},
'grpo': {'adv_estimator': {'minus_baseline': True,
'name': 'grpo',
'normalize_rewards': True,
'use_leave_one_out_baseline': True},
'async_grpo': {'enabled': False,
'in_flight_weight_updates': False,
'max_trajectory_age_steps': 1,
'recompute_kv_cache_after_weight_updates': False},
'batch_multiplier': 1,
'dynamic_sampling_max_gen_batches': 10,
'max_num_epochs': 1,
'max_num_steps': 1000000,
'max_rollout_turns': 1,
'max_val_samples': 256,
'normalize_rewards': True,
'num_generations_per_prompt': 16,
'num_prompts_per_step': 32,
'overlong_filtering': False,
'reward_scaling': {'enabled': False,
'source_max': 1.0,
'source_min': 0.0,
'target_max': 1.0,
'target_min': 0.0},
'reward_shaping': {'enabled': False,
'max_response_length': 512,
'overlong_buffer_length': 128,
'overlong_buffer_penalty': 1,
'stop_properly_penalty_coef': None},
'seed': 42,
'seq_logprob_error_threshold': None,
'use_dynamic_sampling': False,
'use_leave_one_out_baseline': True,
'val_at_end': False,
'val_at_start': False,
'val_batch_size': 256,
'val_period': 10},
'logger': {'gpu_monitoring': {'collection_interval': 10, 'flush_interval': 10},
'log_dir': 'logs',
'mlflow': {'experiment_name': 'grpo-dev',
'run_name': 'sj_megatron_1B',
'tracking_uri': 'http://localhost:5000'},
'mlflow_enabled': False,
'monitor_gpus': False,
'num_val_samples_to_print': 0,
'swanlab': {'name': 'sj_megatron_1B', 'project': 'grpo-dev'},
'swanlab_enabled': False,
'tensorboard': {},
'tensorboard_enabled': False,
'wandb': {'name': 'sj_megatron_1B', 'project': 'grpo-dev'},
'wandb_enabled': False},
'loss_fn': {'force_on_policy_ratio': False,
'kl_input_clamp_value': 20.0,
'kl_output_clamp_value': 10.0,
'ratio_clip_c': None,
'ratio_clip_max': 0.2,
'ratio_clip_min': 0.2,
'reference_policy_kl_penalty': 0.01,
'reference_policy_kl_type': 'k3',
'sequence_level_importance_ratios': False,
'token_level_loss': True,
'truncated_importance_sampling_ratio': None,
'truncated_importance_sampling_ratio_min': None,
'truncated_importance_sampling_type': 'tis',
'use_importance_sampling_correction': False,
'use_kl_in_reward': False,
'use_on_policy_kl_approximation': False},
'policy': {'draft': {'aux_layer_indices': None,
'enabled': False,
'loss_weight': 0.1,
'model_name': None,
'num_layers': None},
'dtensor_cfg': {'_v2': True,
'activation_checkpointing': False,
'automodel_kwargs': {},
'context_parallel_size': 1,
'cpu_offload': False,
'custom_parallel_plan': None,
'enabled': False,
'lora_cfg': {'alpha': 32,
'dim': 8,
'dropout': 0.0,
'dropout_position': 'post',
'enabled': False,
'exclude_modules': [],
'lora_A_init': 'xavier',
'match_all_linear': True,
'target_modules': [],
'use_triton': True},
'sequence_parallel': False,
'tensor_parallel_size': 1},
'dynamic_batching': {'enabled': False,
'logprob_mb_tokens': 512,
'sequence_length_round': 64,
'train_mb_tokens': 512},
'generation': {'backend': 'vllm',
'colocated': {'enabled': True,
'resources': {'gpus_per_node': None,
'num_nodes': None}},
'max_new_tokens': 512,
'mcore_generation_config': {'block_size_tokens': 256,
'buffer_size_gb': 10,
'enable_chunked_prefill': False,
'max_tokens': 16384,
'num_cuda_graphs': 4,
'unified_memory_level': 0,
'use_cuda_graphs_for_non_decode_steps': True},
'stop_strings': None,
'stop_token_ids': None,
'temperature': 1.0,
'top_k': None,
'top_p': 1.0,
'vllm_cfg': {'async_engine': False,
'enable_vllm_metrics_logger': True,
'enforce_eager': False,
'expert_parallel_size': 1,
'gpu_memory_utilization': 0.6,
'kv_cache_dtype': 'auto',
'max_model_len': 512,
'num_first_layers_in_bf16': 0,
'num_last_layers_in_bf16': 0,
'pipeline_parallel_size': 1,
'precision': 'bfloat16',
'tensor_parallel_size': 1,
'use_deep_gemm': False,
'vllm_metrics_logger_interval': 0.5},
'vllm_kwargs': {}},
'generation_batch_size': 64,
'hf_config_overrides': {},
'logprob_batch_size': 1,
'logprob_chunk_size': None,
'make_sequence_length_divisible_by': 1,
'max_grad_norm': 1.0,
'max_total_sequence_length': 512,
'megatron_cfg': {'activation_checkpointing': True,
'apply_rope_fusion': True,
'bias_activation_fusion': True,
'context_parallel_size': 1,
'converter_type': 'Qwen2ForCausalLM',
'defer_fp32_logits': False,
'distributed_data_parallel_config': {'data_parallel_sharding_strategy': 'optim_grads_params',
'grad_reduce_in_fp32': False,
'overlap_grad_reduce': True,
'overlap_param_gather': True,
'use_custom_fsdp': False},
'empty_unused_memory_level': 2,
'enabled': True,
'env_vars': None,
'expert_model_parallel_size': 1,
'expert_tensor_parallel_size': 1,
'force_reconvert_from_hf': False,
'fp8_cfg': {'enabled': False,
'fp8': 'e4m3',
'fp8_param': False,
'fp8_recipe': 'blockwise'},
'freeze_moe_router': True,
'gradient_accumulation_fusion': False,
'moe_enable_deepep': False,
'moe_per_layer_logging': False,
'moe_permute_fusion': False,
'moe_router_bias_update_rate': 0.0,
'moe_router_dtype': 'fp64',
'moe_router_load_balancing_type': 'none',
'moe_shared_expert_overlap': False,
'moe_token_dispatcher_type': 'alltoall',
'num_layers_in_first_pipeline_stage': None,
'num_layers_in_last_pipeline_stage': None,
'optimizer': {'adam_beta1': 0.9,
'adam_beta2': 0.999,
'adam_eps': 1e-08,
'bf16': True,
'clip_grad': 1.0,
'fp16': False,
'lr': 5e-06,
'min_lr': 5e-07,
'optimizer': 'adam',
'optimizer_cpu_offload': False,
'optimizer_offload_fraction': 0.0,
'params_dtype': 'float32',
'sgd_momentum': 0.9,
'use_distributed_optimizer': True,
'use_precision_aware_optimizer': True,
'weight_decay': 0.01},
'peft': {'a2a_experimental': False,
'alpha': 32,
'dim': 8,
'dropout': 0.0,
'dropout_position': 'post',
'enabled': False,
'exclude_modules': [],
'lora_A_init_method': 'xavier',
'lora_B_init_method': 'zero',
'lora_dtype': None,
'target_modules': []},
'pipeline_dtype': 'bfloat16',
'pipeline_model_parallel_size': 1,
'scheduler': {'end_weight_decay': 0.01,
'lr_decay_iters': 1000,
'lr_decay_style': 'constant',
'lr_warmup_init': 5e-07,
'lr_warmup_iters': 13,
'start_weight_decay': 0.01,
'weight_decay_incr_style': 'constant'},
'sequence_parallel': False,
'tensor_model_parallel_size': 1},
'model_name': 'Qwen/Qwen2.5-1.5B',
'offload_optimizer_for_logprob': False,
'optimizer': None,
'precision': 'bfloat16',
'scheduler': [{'kwargs': {'end_factor': 1.0,
'start_factor': 0.1,
'total_iters': 50},
'name': 'torch.optim.lr_scheduler.LinearLR'},
{'kwargs': {'factor': 1.0,
'total_iters': 10000000000},
'name': 'torch.optim.lr_scheduler.ConstantLR'},
{'milestones': [50]}],
'sequence_packing': {'algorithm': 'modified_first_fit_decreasing',
'enabled': True,
'logprob_mb_tokens': 512,
'sequence_length_round': 64,
'train_mb_tokens': 512},
'tokenizer': {'chat_template_kwargs': None,
'name': 'Qwen/Qwen2.5-1.5B'},
'train_global_batch_size': 32,
'train_micro_batch_size': 1}}
📊 Using log directory: logs/exp_026
2026-06-05 14:12:02,406 INFO worker.py:2004 -- Started a local Ray instance. View the dashboard at http://127.0.0.1:8265
/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/_private/worker.py:2052: FutureWarning: Tip: In future versions of Ray, Ray will no longer overrideaccelerator visible devices env var if num_gpus=0 or num_gpus=None (default). To enable this behavior and turn off this error message, set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0
warnings.warn(
No chat template provided, using tokenizer's default

▶ Setting up envs...

▶ Setting up data...
Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 59972.18it/s]
Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 69327.34it/s]
- Loaded training dataset OpenMathInstruct-2 with 950000 samples.
✓ Training dataset loaded with 950000 samples.
- Loaded validation dataset OpenMathInstruct-2 with 50000 samples.
✓ Validation dataset loaded with 50000 samples.
No loggers initialized
✓ Training dataloader loaded with 29687 samples
✓ Validation dataloader loaded with 50000 samples
[WARN] truncated_importance_sampling_ratio is not set, so the following parameters are ignored: truncated_importance_sampling_type. Set truncated_importance_sampling_ratio to enable truncated importance sampling.

▶ Setting up compute cluster...
✓ Ray cluster for policy initialized with 1 nodes

▶ Setting up model and training...
⚙️ Using sequential worker initialization (colocated mode)
Initializing vllm_policy workers: 100%|██████████████████████████████████████████████████████████████████████████████| 1/1 [00:06<00:00, 6.28s/worker]
✓ 1 workers initialized in 6.28s
(VllmGenerationWorker pid=31132) INFO 06-05 14:12:29 [utils.py:238] non-default args: {'served_model_name': 'Qwen/Qwen2.5-1.5B', 'skip_tokenizer_init': True, 'trust_remote_code': True, 'load_format': 'dummy', 'dtype': 'bfloat16', 'max_model_len': 512, 'enable_prefix_caching': True, 'gpu_memory_utilization': 0.6, 'logprobs_mode': 'processed_logprobs', 'worker_extension_cls': 'nemo_rl.models.generation.vllm.vllm_backend.VllmInternalWorkerExtension', 'enable_sleep_mode': True, 'model': 'Qwen/Qwen2.5-1.5B'}
(VllmGenerationWorker pid=31132) WARNING 06-05 14:12:29 [envs.py:1710] Unknown vLLM environment variable detected: VLLM_USE_V1
(VllmGenerationWorker pid=31132) The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
(VllmGenerationWorker pid=31132) WARNING 06-05 14:12:30 [arg_utils.py:1321] The global random seed is set to 0. Since VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may affect the random state of the Python process that launched vLLM.
(VllmGenerationWorker pid=31132) ERROR 06-05 14:12:35 [repo_utils.py:47] Error retrieving file list: [SSL: UNEXPECTED_EOF_WHILE_READING] EOF occurred in violation of protocol (_ssl.c:1032), retrying 1 of 2
(VllmGenerationWorker pid=31132) ERROR 06-05 14:12:42 [repo_utils.py:45] Error retrieving file list: [SSL: UNEXPECTED_EOF_WHILE_READING] EOF occurred in violation of protocol (_ssl.c:1032)
(VllmGenerationWorker pid=31132) ERROR 06-05 14:12:42 [repo_utils.py:110] Error retrieving file list. Please ensure your `model_name_or_path``repo_type`, `token` and `revision` arguments are correctly set. Returning an empty list.
(VllmGenerationWorker pid=31132) The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
(VllmGenerationWorker pid=31132) ERROR 06-05 14:12:57 [repo_utils.py:47] Error retrieving file list: [SSL: UNEXPECTED_EOF_WHILE_READING] EOF occurred in violation of protocol (_ssl.c:1032), retrying 1 of 2
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:02 [model.py:531] Resolved architecture: Qwen2ForCausalLM
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:02 [model.py:1554] Using max model len 512
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:02 [arg_utils.py:1539] Using ray runtime env (env vars redacted): {'py_executable': '/opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker/bin/python', 'env_vars': {'AVAILABLE_ADDR_LIST': '***', 'AVAILABLE_PORT_LIST': '***', 'AWS_OFI_NCCL_VERSION': '***', 'BASH_ENV': '***', 'BROWSER': '***', 'CAL_VERSION': '***', 'COLORTERM': '***', 'COPILOT_DEBUG_NONCE': '***', 'CUBLASMP_VERSION': '***', 'CUBLAS_VERSION': '***', 'CUDA_ARCH_LIST': '***', 'CUDA_DRIVER_VERSION': '***', 'CUDA_VERSION': '***', 'CUDNN_FRONTEND_VERSION': '***', 'CUDNN_HOME': '***', 'CUDNN_VERSION': '***', 'CUFFT_VERSION': '***', 'CUFILE_VERSION': '***', 'CURAND_VERSION': '***', 'CUSOLVER_VERSION': '***', 'CUSPARSELT_VERSION': '***', 'CUSPARSE_VERSION': '***', 'DALI_BUILD': '***', 'DALI_URL_SUFFIX': '***', 'DALI_VERSION': '***', 'DISPLAY': '***', 'EFA_VERSION': '***', 'ENV': '***', 'GDRCOPY_VERSION': '***', 'GIT_ASKPASS': '***', 'HF_ENDPOINT': '***', 'HOME': '***', 'HOSTNAME': '***', 'HPCX_VERSION': '***', 'KMP_DUPLICATE_LIB_OK': '***', 'KMP_INIT_AT_FORK': '***', 'LANG': '***', 'LC_CTYPE': '***', 'LD_LIBRARY_PATH': '***', 'LESSCLOSE': '***', 'LESSOPEN': '***', 'LIBRARY_PATH': '***', 'LOCAL_RANK': '***', 'LS_COLORS': '***', 'MASTER_ADDR': '***', 'MASTER_PORT': '***', 'MODEL_OPT_VERSION': '***', 'MOFED_VERSION': '***', 'NCCL_VERSION': '***', 'NEMO_RL_COMMIT': '***', 'NEMO_RL_VENV_DIR': '***', 'NODE_RANK': '***', 'NPP_VERSION': '***', 'NRL_CONTAINER': '***', 'NSIGHT_COMPUTE_VERSION': '***', 'NSIGHT_SYSTEMS_VERSION': '***', 'NVIDIA_BUILD_ID': '***', 'NVIDIA_BUILD_REF': '***', 'NVIDIA_DRIVER_CAPABILITIES': '***', 'NVIDIA_PRODUCT_NAME': '***', 'NVIDIA_REQUIRE_CUDA': '***', 'NVIDIA_REQUIRE_JETPACK_HOST_MOUNTS': '***', 'NVIDIA_VISIBLE_DEVICES': '***', 'NVJITLINK_VERSION': '***', 'NVJPEG_VERSION': '***', 'NVSHMEM_VERSION': '***', 'OMPI_MCA_coll_hcoll_enable': '***', 'OPAL_PREFIX': '***', 'OPENMPI_VERSION': '***', 'OPENUCX_VERSION': '***', 'PATH': '***', 'POLYGRAPHY_VERSION': '***', 'PWD': '***', 'PYTHONBREAKPOINT': '***', 'RANK': '***', 'RAY_CLIENT_MODE': '***', 'RAY_ENABLE_UV_RUN_RUNTIME_ENV': '***', 'RAY_USAGE_STATS_ENABLED': '***', 'RDMACORE_VERSION': '***', 'REMOTE_CONTAINERS': '***', 'REMOTE_CONTAINERS_DISPLAY_SOCK': '***', 'REMOTE_CONTAINERS_IPC': '***', 'REMOTE_CONTAINERS_SOCKETS': '***', 'SHELL': '***', 'SHLVL': '***', 'SWANLAB_API_HOST': '***', 'SWANLAB_RUNTIME': '***', 'SWANLAB_WEB_HOST': '***', 'TERM': '***', 'TERM_PROGRAM': '***', 'TERM_PROGRAM_VERSION': '***', 'TORCHINDUCTOR_CACHE_DIR': '***', 'TORCH_CUDA_ARCH_LIST': '***', 'TRANSFORMER_ENGINE_VERSION': '***', 'TRTOSS_VERSION': '***', 'TRT_VERSION': '***', 'UV': '***', 'UV_LINK_MODE': '***', 'UV_PROJECT_ENVIRONMENT': '***', 'UV_RUN_RECURSION_DEPTH': '***', 'VIRTUAL_ENV': '***', 'VLLM_CACHE_ROOT': '***', 'VLLM_ENABLE_V1_MULTIPROCESSING': '***', 'VLLM_SKIP_P2P_CHECK': '***', 'VSCODE_GIT_ASKPASS_EXTRA_ARGS': '***', 'VSCODE_GIT_ASKPASS_MAIN': '***', 'VSCODE_GIT_ASKPASS_NODE': '***', 'VSCODE_GIT_IPC_HANDLE': '***', 'VSCODE_IPC_HOOK_CLI': '***', 'VSCODE_PYTHON_AUTOACTIVATE_GUARD': '***', 'WORLD_SIZE': '***', '_': '***', '_CUDA_COMPAT_PATH': '***', '_MLFLOW_TELEMETRY_SESSION_ID': '***'}}
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:02 [scheduler.py:231] Chunked prefill is enabled with max_num_batched_tokens=8192.
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:02 [vllm.py:747] Asynchronous scheduling is enabled.
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:04 [core.py:101] Initializing a V1 LLM engine (v0.17.1) with config: model='Qwen/Qwen2.5-1.5B', speculative_config=None, tokenizer='Qwen/Qwen2.5-1.5B', skip_tokenizer_init=True, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=512, download_dir=None, load_format=dummy, tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=cuda, structured_outputs_config=StructuredOutputsConfig(backend='auto', disable_fallback=False, disable_any_whitespace=False, disable_additional_properties=False, reasoning_parser='', reasoning_parser_plugin='', enable_in_reasoning=False), observability_config=ObservabilityConfig(show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None, kv_cache_metrics=False, kv_cache_metrics_sample=0.01, cudagraph_metrics=False, enable_layerwise_nvtx_tracing=False, enable_mfu_metrics=False, enable_mm_processor_stats=False, enable_logging_iteration_details=False), seed=0, served_model_name=Qwen/Qwen2.5-1.5B, enable_prefix_caching=True, enable_chunked_prefill=True, pooler_config=None, compilation_config={'level': None, 'mode': , 'debug_dump_path': None, 'cache_dir': '', 'compile_cache_save_format': 'binary', 'backend': 'inductor', 'custom_ops': ['none'], 'splitting_ops': ['vllm::unified_attention', 'vllm::unified_attention_with_output', 'vllm::unified_mla_attention', 'vllm::unified_mla_attention_with_output', 'vllm::mamba_mixer2', 'vllm::mamba_mixer', 'vllm::short_conv', 'vllm::linear_attention', 'vllm::plamo2_mamba_mixer', 'vllm::gdn_attention_core', 'vllm::kda_attention', 'vllm::sparse_attn_indexer', 'vllm::rocm_aiter_sparse_attn_indexer', 'vllm::unified_kv_cache_update', 'vllm::unified_mla_kv_cache_update'], 'compile_mm_encoder': False, 'compile_sizes': [], 'compile_ranges_split_points': [8192], 'inductor_compile_config': {'enable_auto_functionalized_v2': False, 'combo_kernels': True, 'benchmark_combo_kernel': True}, 'inductor_passes': {}, 'cudagraph_mode': , 'cudagraph_num_of_warmups': 1, 'cudagraph_capture_sizes': [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512], 'cudagraph_copy_inputs': False, 'cudagraph_specialize_lora': True, 'use_inductor_graph_partition': False, 'pass_config': {'fuse_norm_quant': False, 'fuse_act_quant': False, 'fuse_attn_quant': False, 'enable_sp': False, 'fuse_gemm_comms': False, 'fuse_allreduce_rms': False}, 'max_cudagraph_capture_size': 512, 'dynamic_shapes_config': {'type': , 'evaluate_guards': False, 'assume_32_bit_indexing':False}, 'local_cache_dir': None, 'fast_moe_cold_start': True, 'static_all_moe_layers': []}
(VllmGenerationWorker pid=31132) WARNING 06-05 14:13:04 [interface.py:472] Using 'pin_memory=False' as WSL is detected. This may slow down the performance.
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:04 [worker_base.py:283] Injected into for extended collective_rpc calls ['_load_draft_weights', '_maybe_process_fp8_kv_cache','_split_policy_and_draft_weights', 'cleanup', 'get_zmq_address', 'init_collective', 'maybe_init_zmq', 'prepare_refit_info', 'report_device_id', 'start_gpu_profiling', 'stop_gpu_profiling', 'update_weights_from_collective', 'update_weights_via_ipc_zmq']
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:04 [parallel_state.py:1393] world_size=1 rank=0 local_rank=0 distributed_init_method=tcp://172.17.0.2:39531 backend=nccl
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:04 [parallel_state.py:1715] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, PCP rank 0, TP rank 0, EP rank N/A, EPLB rank N/A
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:05 [base.py:106] Offloader set to NoopOffloader
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:05 [gpu_model_runner.py:4281] Starting to load model Qwen/Qwen2.5-1.5B...
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:06 [cuda.py:405] Using FLASH_ATTN attention backend out of potential backends: ['FLASH_ATTN', 'FLASHINFER', 'TRITON_ATTN', 'FLEX_ATTENTION'].
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:06 [flash_attn.py:587] Using FlashAttention version 2
(VllmGenerationWorker pid=31132) :1325: FutureWarning: The cuda.cudart module is deprecated and will be removed in a future release, please switch to use the cuda.bindings.runtime module instead.
(VllmGenerationWorker pid=31132) :1325: FutureWarning: The cuda.nvrtc module is deprecated and will be removed in a future release, please switch to use the cuda.bindings.nvrtc module instead.
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:06 [gpu_model_runner.py:4364] Model loading took 2.91 GiB memory and 0.922607 seconds
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:07 [decorators.py:465] Directly load AOT compilation from path /root/.cache/vllm_0/torch_compile_cache/torch_aot_compile/26ced3911d4d2d243adc39d531ad59a3ccf6bba77bb0344bdbc6933a5cb4725d/rank_0_0/model
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:08 [backends.py:916] Using cache directory: /root/.cache/vllm_0/torch_compile_cache/852c442d25/rank_0_0/backbone for vLLM's torch.compile
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:08 [backends.py:976] Dynamo bytecode transform time: 1.91 s
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:09 [backends.py:266] Directly load the compiled graph(s) for compile range (1, 8192) from the cache,took 0.657 s
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:09 [monitor.py:35] torch.compile takes 2.70 s in total
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:10 [gpu_worker.py:424] Available KV cache memory: 10.7 GiB
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:10 [kv_cache_utils.py:1314] GPU KV cache size: 400,560 tokens
(VllmGenerationWorker pid=31132) INFO 06-05 14:13:10 [kv_cache_utils.py:1319] Maximum concurrency for 512 tokens per request: 782.34x
Capturing CUDA graphs (mixed prefill-decode, PIECEWISE): 0%| | 0/51 [00:00 number of parameters on (tensor, pipeline) model parallel rank (0, 0): 1543714304
(MegatronPolicyWorker pid=35477) saving checkpoint at iteration 0 to /root/.cache/huggingface/nemo_rl/Qwen/Qwen2.5-1.5B in torch_dist format
(MegatronPolicyWorker pid=35477) WARNING:megatron.core.rerun_state_machine:_rank_utils.py:43: Implicit initialization of Rerun State Machine!
(MegatronPolicyWorker pid=35477) WARNING:megatron.core.rerun_state_machine:_rank_utils.py:43: RerunStateMachine initialized in mode RerunMode.DISABLED
(MegatronPolicyWorker pid=35477) Exception raised in creation task: The actor died because of an error raised in its creation task, ray::lm_policy-0-0:MegatronPolicyWorker.__init__() (pid=35477, ip=172.17.0.2, actor_id=362b2a9f9ed237b34fcb2fd901000000, repr=MegatronPolicyWorker[rank=0])
(MegatronPolicyWorker pid=35477) File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 140, in __init__
(MegatronPolicyWorker pid=35477) handle_model_import(
(MegatronPolicyWorker pid=35477) ~~~~~~~~~~~~~~~~~~~^
(MegatronPolicyWorker pid=35477) config, hf_model_name, pretrained_path, pt_checkpoint_exists
(MegatronPolicyWorker pid=35477) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) )
(MegatronPolicyWorker pid=35477) ^
(MegatronPolicyWorker pid=35477) File "/opt/nemo-rl/nemo_rl/models/megatron/setup.py", line 968, in handle_model_import
(MegatronPolicyWorker pid=35477) import_model_from_hf_name(
(MegatronPolicyWorker pid=35477) ~~~~~~~~~~~~~~~~~~~~~~~~~^
(MegatronPolicyWorker pid=35477) hf_model_name,
(MegatronPolicyWorker pid=35477) ^^^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) ...<2 lines>...
(MegatronPolicyWorker pid=35477) **hf_config_overrides,
(MegatronPolicyWorker pid=35477) ^^^^^^^^^^^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) )
(MegatronPolicyWorker pid=35477) ^
(MegatronPolicyWorker pid=35477) File "/opt/nemo-rl/nemo_rl/models/megatron/community_import.py", line 107, in import_model_from_hf_name
(MegatronPolicyWorker pid=35477) bridge.save_megatron_model(megatron_model, output_path)
(MegatronPolicyWorker pid=35477) ~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) File "/opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src/megatron/bridge/models/conversion/auto_bridge.py", line 939, in save_megatron_model
(MegatronPolicyWorker pid=35477) save_megatron_model(
(MegatronPolicyWorker pid=35477) ~~~~~~~~~~~~~~~~~~~^
(MegatronPolicyWorker pid=35477) model,
(MegatronPolicyWorker pid=35477) ^^^^^^
(MegatronPolicyWorker pid=35477) ...<3 lines>...
(MegatronPolicyWorker pid=35477) hf_tokenizer_kwargs=hf_tokenizer_kwargs,
(MegatronPolicyWorker pid=35477) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) )
(MegatronPolicyWorker pid=35477) ^
(MegatronPolicyWorker pid=35477) File "/opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src/megatron/bridge/training/model_load_save.py", line 711, in save_megatron_model
(MegatronPolicyWorker pid=35477) save_checkpoint(
(MegatronPolicyWorker pid=35477) ~~~~~~~~~~~~~~~^
(MegatronPolicyWorker pid=35477) state=state,
(MegatronPolicyWorker pid=35477) ^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) ...<4 lines>...
(MegatronPolicyWorker pid=35477) callback_manager=None,
(MegatronPolicyWorker pid=35477) ^^^^^^^^^^^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) )
(MegatronPolicyWorker pid=35477) ^
(MegatronPolicyWorker pid=35477) File "/opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src/megatron/bridge/training/checkpointing.py", line 947, in save_checkpoint
(MegatronPolicyWorker pid=35477) async_save_request = dist_checkpointing.save(
(MegatronPolicyWorker pid=35477) state_dict,
(MegatronPolicyWorker pid=35477) ...<6 lines>...
(MegatronPolicyWorker pid=35477) async_strategy=ckpt_cfg.async_strategy,
(MegatronPolicyWorker pid=35477) )
(MegatronPolicyWorker pid=35477) File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 390, in save
(MegatronPolicyWorker pid=35477) sharded_strategy.save(sharded_state_dict, checkpoint_dir)
(MegatronPolicyWorker pid=35477) ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/dist_checkpointing/strategies/fully_parallel.py", line 103, in save
(MegatronPolicyWorker pid=35477) return self.base_strategy.save(sharded_state_dict, checkpoint_dir)
(MegatronPolicyWorker pid=35477) ~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/dist_checkpointing/strategies/torch.py", line 650, in save
(MegatronPolicyWorker pid=35477) async_request.execute_sync()
(MegatronPolicyWorker pid=35477) ~~~~~~~~~~~~~~~~~~~~~~~~~~^^
(MegatronPolicyWorker pid=35477) File "/opt/ray_venvs/nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker/lib/python3.13/site-packages/nvidia_resiliency_ext/checkpointing/async_ckpt/core.py", line 164, in execute_sync
(MegatronPolicyWorker pid=35477) async_fn_args[1] = self.preload_fn()
(MegatronPolicyWorker pid=35477) ~~~~~~~~~~~~~~~^^
(MegatronPolicyWorker pid=35477) File "/opt/ray_venvs/nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker/lib/python3.13/site-packages/nvidia_resiliency_ext/checkpointing/async_ckpt/filesystem_async.py", line 527, in preload_tensors
(MegatronPolicyWorker pid=35477) tensor_list.append((item, tensor.to("cpu", non_blocking=non_blocking)))
(MegatronPolicyWorker pid=35477) ~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(MegatronPolicyWorker pid=35477) torch.AcceleratorError: CUDA error: out of memory
(MegatronPolicyWorker pid=35477) Search for `cudaErrorMemoryAllocation' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html formore information.
(MegatronPolicyWorker pid=35477) CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
(MegatronPolicyWorker pid=35477) For debugging consider passing CUDA_LAUNCH_BLOCKING=1
(MegatronPolicyWorker pid=35477) Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
Traceback (most recent call last):
File "/opt/nemo-rl/examples/run_grpo.py", line 187, in
main()
~~~~^^
File "/opt/nemo-rl/examples/run_grpo.py", line 113, in main
) = setup(config, tokenizer, dataset, val_dataset)
~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/nemo_rl/algorithms/grpo.py", line 724, in setup
policy.print_node_ip_and_gpu_id()
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
File "/opt/nemo-rl/nemo_rl/models/policy/lm_policy.py", line 997, in print_node_ip_and_gpu_id
results = ray.get(
self.worker_group.run_all_workers_single_data(
"report_node_ip_and_gpu_id",
)
)
File "/opt/nemo_rl_venv/lib/python3.13/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.13/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
return func(*args, **kwargs)
File "/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/_private/worker.py", line 2981, in get
values, debugger_breakpoint = worker.get_objects(
~~~~~~~~~~~~~~~~~~^
object_refs, timeout, use_object_store=_use_object_store
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/_private/worker.py", line 1014, in get_objects
raise value
ray.exceptions.ActorDiedError: The actor died because of an error raised in its creation task, ray::lm_policy-0-0:MegatronPolicyWorker.__init__() (pid=35477, ip=172.17.0.2, actor_id=362b2a9f9ed237b34fcb2fd901000000, repr=MegatronPolicyWorker[rank=0])
File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 140, in __init__
handle_model_import(
~~~~~~~~~~~~~~~~~~~^
config, hf_model_name, pretrained_path, pt_checkpoint_exists
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/nemo-rl/nemo_rl/models/megatron/setup.py", line 968, in handle_model_import
import_model_from_hf_name(
~~~~~~~~~~~~~~~~~~~~~~~~~^
hf_model_name,
^^^^^^^^^^^^^^
...<2 lines>...
**hf_config_overrides,
^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/nemo-rl/nemo_rl/models/megatron/community_import.py", line 107, in import_model_from_hf_name
bridge.save_megatron_model(megatron_model, output_path)
~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src/megatron/bridge/models/conversion/auto_bridge.py", line 939, in save_megatron_model
save_megatron_model(
~~~~~~~~~~~~~~~~~~~^
model,
^^^^^^
...<3 lines>...
hf_tokenizer_kwargs=hf_tokenizer_kwargs,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src/megatron/bridge/training/model_load_save.py", line 711, in save_megatron_model
save_checkpoint(
~~~~~~~~~~~~~~~^
state=state,
^^^^^^^^^^^^
...<4 lines>...
callback_manager=None,
^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src/megatron/bridge/training/checkpointing.py", line 947, in save_checkpoint
async_save_request = dist_checkpointing.save(
state_dict,
...<6 lines>...
async_strategy=ckpt_cfg.async_strategy,
)
File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 390, in save
sharded_strategy.save(sharded_state_dict, checkpoint_dir)
~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/dist_checkpointing/strategies/fully_parallel.py", line 103, in save
return self.base_strategy.save(sharded_state_dict, checkpoint_dir)
~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/dist_checkpointing/strategies/torch.py", line 650, in save
async_request.execute_sync()
~~~~~~~~~~~~~~~~~~~~~~~~~~^^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker/lib/python3.13/site-packages/nvidia_resiliency_ext/checkpointing/async_ckpt/core.py", line 164, in execute_sync
async_fn_args[1] = self.preload_fn()
~~~~~~~~~~~~~~~^^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker/lib/python3.13/site-packages/nvidia_resiliency_ext/checkpointing/async_ckpt/filesystem_async.py", line 527, in preload_tensors
tensor_list.append((item, tensor.to("cpu", non_blocking=non_blocking)))
~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.AcceleratorError: CUDA error: out of memory
Search for `cudaErrorMemoryAllocation' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
Error removing placement group PlacementGroupID(69f9c6c23fb4ab33261bd5d5192e01000000): sys.meta_path is None, Python is likely shutting down
Error during policy shutdown: sys.meta_path is None, Python is likely shutting down
Exception ignored in:
Traceback (most recent call last):
File "/opt/nemo-rl/nemo_rl/models/policy/lm_policy.py", line 983, in __del__
File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 1006, in shutdown
File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 817, in run_all_workers_single_data
File "/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/actor.py", line 676, in remote
File "/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/_private/auto_init_hook.py", line 21, in auto_init_wrapper
File "/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/_private/auto_init_hook.py", line 15, in auto_init_ray
File "/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
File "/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/_private/worker.py", line 1757, in init
File "/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/job_config.py", line 72, in __init__
File "/opt/nemo_rl_venv/lib/python3.13/site-packages/ray/job_config.py", line 155, in set_default_actor_lifetime
ImportError: sys.meta_path is None, Python is likely shutting down