diff --git a/docs/design-docs/generation.md b/docs/design-docs/generation.md index 87c40084a58..803603cc106 100644 --- a/docs/design-docs/generation.md +++ b/docs/design-docs/generation.md @@ -247,9 +247,8 @@ policy: max_tokens: 16384 ``` -For a complete example, see: -- **Configuration**: `examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml` -- **Test Script**: `tests/functional/grpo_megatron_generation.sh` +For a complete example, see +`examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml`. ## Extend with New Backends diff --git a/docs/guides/async-grpo.md b/docs/guides/async-grpo.md index 6de30e72691..fa54b19fb9a 100644 --- a/docs/guides/async-grpo.md +++ b/docs/guides/async-grpo.md @@ -24,7 +24,7 @@ loss_fn: use_importance_sampling_correction: true ``` -3. **Disable colocated inference** (required for async mode with the vLLM backend; the Megatron backend supports colocated async — see `examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml`): +3. **Disable colocated inference** (required for async mode with the vLLM backend; the Megatron backend supports colocated async — see `examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml`): ```yaml policy: generation: diff --git a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated.yaml b/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated.yaml deleted file mode 100644 index 9e9577adb18..00000000000 --- a/examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated.yaml +++ /dev/null @@ -1,41 +0,0 @@ -defaults: ../../grpo_math_1B.yaml -grpo: - max_num_steps: 30 -checkpointing: - enabled: false - checkpoint_dir: results/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated - save_period: 100 -policy: - model_name: meta-llama/Llama-3.2-1B-Instruct - tokenizer: - name: meta-llama/Llama-3.2-1B-Instruct - optimizer: null - megatron_cfg: - enabled: true - scheduler: - lr_warmup_iters: 50 - dtensor_cfg: - enabled: false - make_sequence_length_divisible_by: 1 - generation: - backend: megatron - max_new_tokens: 512 - vllm_cfg: - max_model_len: 512 - colocated: - enabled: false - resources: - gpus_per_node: 8 - num_nodes: 1 -data: - max_input_seq_length: 512 -logger: - log_dir: logs/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated - wandb_enabled: true - tensorboard_enabled: true - wandb: - project: nemo-rl - name: grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated -cluster: - gpus_per_node: 8 - num_nodes: 2 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick.yaml new file mode 100644 index 00000000000..cb84214fccd --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick.yaml @@ -0,0 +1,9 @@ +defaults: ./grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml +# Reduced mcore-lane smoke leg: 3 steps, no checkpointing, gen_kl_error-only +# gate (see the suite). Exists so Megatron-Core/Bridge bumps exercise the 30B +# megatron-inference stack without paying for the full nightly leg. +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick +logger: + wandb: + name: grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml new file mode 100644 index 00000000000..ead10f761f9 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml @@ -0,0 +1,36 @@ +defaults: ../../../nemo_gym/grpo_nanov3.yaml +# Async colocated GRPO + gym on Megatron generation with a resharded inference +# layout: training runs TP2xCP2xEP8 across all 16 GPUs; every wake reshards the +# weights into TP4xEP4 inference_optimized models on the same GPUs. One nightly +# covers the async loop with in-flight updates, the gym HTTP ingress, colocated +# reshard, and (via the suite's save_period=val_period=4) the save-bound +# deferred-wake path. +grpo: + async_grpo: + enabled: true + in_flight_weight_updates: true + overlong_filtering: false +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym +policy: + model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 + tokenizer: + name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + train_global_batch_size: 16 + max_total_sequence_length: 2048 + megatron_cfg: + pipeline_model_parallel_size: 1 + context_parallel_size: 2 + generation: + backend: megatron + mcore_generation_config: + tensor_model_parallel_size: 4 + expert_model_parallel_size: 4 +logger: + wandb_enabled: true + tensorboard_enabled: true + wandb: + project: nemo-rl + name: grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym +cluster: + num_nodes: 2 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick.yaml new file mode 100644 index 00000000000..1a6382ac8c0 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick.yaml @@ -0,0 +1,9 @@ +defaults: ./grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml +# Reduced mcore-lane smoke leg: 3 steps, no checkpointing, gen_kl_error-only +# gate (see the suite). Exists so Megatron-Core/Bridge bumps exercise the 30B +# megatron-inference stack without paying for the full nightly leg. +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick +logger: + wandb: + name: grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml similarity index 81% rename from examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.yaml rename to examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml index c2ba70d7b72..f171bcd913f 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml @@ -5,7 +5,7 @@ grpo: in_flight_weight_updates: true overlong_filtering: false checkpointing: - checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym + checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym policy: model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 tokenizer: @@ -17,6 +17,8 @@ policy: context_parallel_size: 2 generation: backend: megatron + mcore_generation_config: + expert_model_parallel_size: 4 colocated: enabled: false resources: @@ -27,6 +29,6 @@ logger: tensorboard_enabled: true wandb: project: nemo-rl - name: grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym + name: grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym cluster: num_nodes: 2 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml deleted file mode 100644 index 69abf1b560d..00000000000 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml +++ /dev/null @@ -1,50 +0,0 @@ -defaults: ../../grpo_math_1B.yaml -# Async colocated GRPO with Megatron (Nemotron-3-Nano-30B-A3B): training and -# generation share GPUs/workers; the engine sleeps across training steps and -# serves the shared weights (resharding into a dedicated inference layout when -# one is configured). Off-policy data via the replay buffer. -grpo: - num_prompts_per_step: 2 - num_generations_per_prompt: 8 - async_grpo: - enabled: true - max_trajectory_age_steps: 4 # weight versions a rollout may span - in_flight_weight_updates: true - # For AREAL-style KV invalidation, add recompute_kv_cache_after_weight_updates: true. -loss_fn: - use_importance_sampling_correction: true # required for off-policy replay data -checkpointing: - enabled: false - checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated - save_period: 100 -policy: - model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 - tokenizer: - name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - train_global_batch_size: 16 - train_micro_batch_size: 1 - logprob_batch_size: 1 - max_total_sequence_length: 2048 - megatron_cfg: - enabled: true - bias_activation_fusion: false - tensor_model_parallel_size: 2 - expert_model_parallel_size: 8 - sequence_parallel: true - dtensor_cfg: - enabled: false - sequence_packing: - enabled: false - generation: - backend: megatron -logger: - log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated - wandb_enabled: true - tensorboard_enabled: true - wandb: - project: nemo-rl - name: grpo-nanov3-30BA3B-4n4g-megatron_async_colocated -cluster: - gpus_per_node: 4 - num_nodes: 4 - segment_size: 2 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml deleted file mode 100644 index 3c45dc1c8b3..00000000000 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml +++ /dev/null @@ -1,47 +0,0 @@ -defaults: ../../grpo_math_1B.yaml -grpo: - num_prompts_per_step: 2 - num_generations_per_prompt: 8 - max_num_steps: 500 -checkpointing: - enabled: false - checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard - save_period: 100 -policy: - model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 - tokenizer: - name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - train_global_batch_size: 16 - train_micro_batch_size: 1 - logprob_batch_size: 1 - max_total_sequence_length: 2048 - megatron_cfg: - enabled: true - bias_activation_fusion: false - tensor_model_parallel_size: 2 - expert_model_parallel_size: 8 - sequence_parallel: true - dtensor_cfg: - enabled: false - sequence_packing: - enabled: false - generation: - backend: megatron - mcore_generation_config: - transformer_impl: inference_optimized - moe_router_dtype: fp32 - activation_checkpointing: false - tensor_model_parallel_size: 4 - expert_model_parallel_size: 4 - sequence_parallel: true -logger: - log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard - wandb_enabled: true - tensorboard_enabled: true - wandb: - project: nemo-rl - name: grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard -cluster: - gpus_per_node: 4 - num_nodes: 4 - segment_size: 2 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick.yaml new file mode 100644 index 00000000000..df6ee1255e5 --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick.yaml @@ -0,0 +1,9 @@ +defaults: ./grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.yaml +# Reduced mcore-lane smoke leg: 3 steps, no checkpointing, gen_kl_error-only +# gate (see the suite). Exists so Megatron-Core/Bridge bumps exercise the 30B +# megatron-inference stack without paying for the full nightly leg. +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick +logger: + wandb: + name: grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.yaml new file mode 100644 index 00000000000..96b9a52263a --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.yaml @@ -0,0 +1,14 @@ +defaults: ./grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml +# GB200 sizing: the same 16-GPU layouts on 4-GPU nodes, so the training EP8 +# groups span two nodes (cross-node expert parallelism, the production +# Blackwell shape) and placement runs through the NVLink-domain topology path; +# inference TP4 occupies exactly one node. +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym +logger: + wandb: + name: grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym +cluster: + gpus_per_node: 4 + num_nodes: 4 + segment_size: 2 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick.yaml new file mode 100644 index 00000000000..22a351449fe --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick.yaml @@ -0,0 +1,9 @@ +defaults: ./grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.yaml +# Reduced mcore-lane smoke leg: 3 steps, no checkpointing, gen_kl_error-only +# gate (see the suite). Exists so Megatron-Core/Bridge bumps exercise the 30B +# megatron-inference stack without paying for the full nightly leg. +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick +logger: + wandb: + name: grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.yaml new file mode 100644 index 00000000000..fdeba1ac08c --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.yaml @@ -0,0 +1,19 @@ +defaults: ./grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml +# GB200 sizing: the same 8+8-GPU non-colocated split on 4-GPU nodes, so the +# training group (CP2xEP8) and the dedicated inference group each span two +# nodes and placement runs through the NVLink-domain topology path. +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym +policy: + generation: + colocated: + resources: + gpus_per_node: 4 + num_nodes: 2 +logger: + wandb: + name: grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym +cluster: + gpus_per_node: 4 + num_nodes: 4 + segment_size: 2 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml similarity index 56% rename from examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.yaml rename to examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml index ed4a46947bf..1d78a49b1d1 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml @@ -1,16 +1,19 @@ defaults: ../../grpo_math_1B.yaml +# BF16 training with MXFP8-quantized Megatron inference on a dedicated 8-GPU group: +# every weight sync quantizes BF16 -> MXFP8 over the NCCL refit, +# with EP8-sharded expert transfers across NVLink domains. grpo: num_prompts_per_step: 2 num_generations_per_prompt: 8 checkpointing: - checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation + enabled: false + checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts policy: model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 tokenizer: name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 train_global_batch_size: 16 train_micro_batch_size: 1 - logprob_batch_size: 1 max_total_sequence_length: 2048 dtensor_cfg: enabled: false @@ -27,24 +30,29 @@ policy: backend: megatron mcore_generation_config: transformer_impl: inference_optimized - inference_grouped_gemm_backend: "vllm" - moe_router_num_groups: null - moe_router_group_topk: null - activation_checkpointing: false + inference_grouped_gemm_backend: "torch" + inference_moe_token_dispatcher_type: nvls tensor_model_parallel_size: 2 expert_model_parallel_size: 8 sequence_parallel: true + logprobs_mode: raw_logprobs + fp8_cfg: + enabled: true + fp8: e4m3 + fp8_recipe: mxfp8 + fp8_param: true colocated: enabled: false resources: - gpus_per_node: 8 - num_nodes: 1 + gpus_per_node: 4 + num_nodes: 2 logger: wandb_enabled: true tensorboard_enabled: true wandb: project: nemo-rl - name: grpo-nanov3-30BA3B-2n8g-megatron_generation + name: grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts cluster: - gpus_per_node: 8 - num_nodes: 2 + gpus_per_node: 4 + num_nodes: 4 + segment_size: 2 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.yaml deleted file mode 100644 index 78f43ccb3c5..00000000000 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.yaml +++ /dev/null @@ -1,41 +0,0 @@ -defaults: ../../grpo_math_1B.yaml -grpo: - num_prompts_per_step: 2 - num_generations_per_prompt: 8 -loss_fn: - use_importance_sampling_correction: true # matches the async sibling so the pair isolates the loop -checkpointing: - enabled: false - checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated - save_period: 100 -policy: - model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16 - tokenizer: - name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - train_global_batch_size: 16 - train_micro_batch_size: 1 - logprob_batch_size: 1 - max_total_sequence_length: 2048 - megatron_cfg: - enabled: true - bias_activation_fusion: false - tensor_model_parallel_size: 2 - expert_model_parallel_size: 8 - sequence_parallel: true - dtensor_cfg: - enabled: false - sequence_packing: - enabled: false - generation: - backend: megatron -logger: - log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated - wandb_enabled: true - tensorboard_enabled: true - wandb: - project: nemo-rl - name: grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated -cluster: - gpus_per_node: 4 - num_nodes: 4 - segment_size: 2 diff --git a/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.yaml b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.yaml new file mode 100644 index 00000000000..85f5b62ca3d --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.yaml @@ -0,0 +1,12 @@ +defaults: ./grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron-single-controller-sync.yaml +logger: + log_dir: logs/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync + wandb: + name: grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync +checkpointing: + checkpoint_dir: results/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync +policy: + generation: + backend: megatron + mcore_generation_config: + kv_cache_management_mode: recompute diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml index 0f741856edf..eb674c74ffa 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml @@ -75,3 +75,4 @@ logger: cluster: gpus_per_node: 4 num_nodes: 8 + segment_size: 2 diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index d5d5ec58233..228c85c0c9b 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1817,7 +1817,6 @@ async def _train_pump(self) -> None: ) if groups_dispatched == 0 and self._gen is not None: - # Raise here for observability. try: await asyncio.to_thread(self._gen.snapshot_step_metrics) except RayActorError as error: diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 737c89664ec..590a8cb44f1 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -777,6 +777,7 @@ def _build_sampling_params( termination_id=self.megatron_tokenizer.eod, stop_words=stop_words, return_prompt_tokens=return_prompt_tokens, + detokenize_stop_sequence=True, ) def _merge_stop_strings( diff --git a/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh b/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh index 6c843738c1e..34910911f33 100644 --- a/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh +++ b/tests/functional/L1_Functional_Tests_GB200_MXFP8.sh @@ -35,9 +35,6 @@ run_test() { } run_test uv run --no-sync bash ./tests/functional/grpo_vllm_mxfp8_rollout_gb200.sh -# Disabled: sized for 4 GPUs (2 train + 2 inference, EP=2 per side) but this shard's -# gcp-gpu-x2 runner has 2, so the inference placement group can never be satisfied (#3731). -# run_test uv run --no-sync bash ./tests/functional/grpo_megatron_mxfp8_refit_gb200.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/L1_Functional_Tests_Megatron_4.sh b/tests/functional/L1_Functional_Tests_Megatron_4.sh index a3e71a6215b..4c0accc98f1 100644 --- a/tests/functional/L1_Functional_Tests_Megatron_4.sh +++ b/tests/functional/L1_Functional_Tests_Megatron_4.sh @@ -34,34 +34,17 @@ run_test() { fi } -# Megatron Inference currently hits an IMA on Blackwell tests. -# TODO: remove this guard once the upstream dependency is bumped. -megatron_generation_supported() { - if command -v nvidia-smi &> /dev/null; then - local compute_cap - compute_cap=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0 2>/dev/null | tr -d '. ' || true) - if [[ "${compute_cap:-0}" -ge 100 ]]; then - echo "WARNING: Skipping Blackwell x Megatron Inference tests" - return 1 - fi - fi - return 0 -} - -if megatron_generation_supported; then - run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation.sh - run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_topology.sh - run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated.sh - run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard.sh - run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_async_grpo.sh - run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh - run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_async_gym.sh - run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_topp_topk.sh - # Disabled: token_mult_prob_error ~2.0 > 1.1 under top_p/top_k after the - # Megatron-LM cf2f07d7 -> bacd3404 bump; see #3385. - # run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_multiturn.sh - run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_gym.sh -fi +run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated.sh +run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated.sh +run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard.sh +run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_gym.sh +run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated_gym.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_topp_topk.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated_async_gym.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard_async_gym.sh +# Disabled: token_mult_prob_error ~2.0 > 1.1 under top_p/top_k after the +# Megatron-LM cf2f07d7 -> bacd3404 bump; see #3385. +# run_test uv run --no-sync bash ./tests/functional/grpo_megatron_generation_multiturn.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/grpo_megatron_generation.sh b/tests/functional/grpo_megatron_generation.sh deleted file mode 100644 index c5a82781b3e..00000000000 --- a/tests/functional/grpo_megatron_generation.sh +++ /dev/null @@ -1,46 +0,0 @@ -#!/bin/bash - -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) -# Mark the current repo as safe, since wandb fetches metadata about the repo -git config --global --add safe.directory $PROJECT_ROOT - -set -eou pipefail - -EXP_NAME=$(basename $0 .sh) -EXP_DIR=$SCRIPT_DIR/$EXP_NAME -LOG_DIR=$EXP_DIR/logs -JSON_METRICS=$EXP_DIR/metrics.json -RUN_LOG=$EXP_DIR/run.log -export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} - -rm -rf $EXP_DIR $LOG_DIR -mkdir -p $EXP_DIR $LOG_DIR - -# Using Qwen2.5-0.5B instead of Qwen3-0.6B because the latter is not supported by Megatron yet -cd $PROJECT_ROOT -uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ - $PROJECT_ROOT/examples/run_grpo.py \ - --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \ - policy.model_name=Qwen/Qwen2.5-0.5B \ - grpo.num_prompts_per_step=2 \ - grpo.num_generations_per_prompt=4 \ - policy.train_global_batch_size=4 \ - policy.logprob_batch_size=4 \ - policy.train_micro_batch_size=1 \ - policy.generation.backend=megatron \ - cluster.gpus_per_node=2 \ - grpo.max_num_steps=2 \ - logger.tensorboard_enabled=true \ - logger.log_dir=$LOG_DIR \ - logger.wandb_enabled=false \ - logger.monitor_gpus=true \ - checkpointing.enabled=false \ - $@ \ - 2>&1 | tee $RUN_LOG - -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS - -uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/token_mult_prob_error"]) < 1.05' - diff --git a/tests/functional/grpo_megatron_generation_topology.sh b/tests/functional/grpo_megatron_generation_colocated.sh similarity index 100% rename from tests/functional/grpo_megatron_generation_topology.sh rename to tests/functional/grpo_megatron_generation_colocated.sh diff --git a/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh b/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh deleted file mode 100644 index f58d0138f5b..00000000000 --- a/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh +++ /dev/null @@ -1,73 +0,0 @@ -#!/bin/bash - -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) -# Mark the current repo as safe, since wandb fetches metadata about the repo -git config --global --add safe.directory $PROJECT_ROOT - -set -eou pipefail - -EXP_NAME=$(basename $0 .sh) -EXP_DIR=$SCRIPT_DIR/$EXP_NAME -LOG_DIR=$EXP_DIR/logs -JSON_METRICS=$EXP_DIR/metrics.json -RUN_LOG=$EXP_DIR/run.log -CKPT_DIR=$EXP_DIR/ckpts -export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} - -rm -rf $EXP_DIR $LOG_DIR -mkdir -p $EXP_DIR $LOG_DIR $CKPT_DIR - -# clean up checkpoint directory on exit -trap "rm -rf $CKPT_DIR" EXIT - -# async colocated -cd $PROJECT_ROOT -uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ - $PROJECT_ROOT/examples/run_grpo.py \ - --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \ - policy.model_name=Qwen/Qwen2.5-0.5B \ - grpo.num_prompts_per_step=2 \ - grpo.num_generations_per_prompt=4 \ - policy.train_global_batch_size=4 \ - policy.logprob_batch_size=4 \ - policy.train_micro_batch_size=1 \ - policy.generation.backend=megatron \ - grpo.async_grpo.enabled=true \ - grpo.async_grpo.max_trajectory_age_steps=1 \ - grpo.async_grpo.in_flight_weight_updates=true \ - loss_fn.use_importance_sampling_correction=true \ - grpo.max_num_steps=3 \ - grpo.val_period=3 \ - grpo.max_val_samples=8 \ - grpo.val_batch_size=8 \ - cluster.gpus_per_node=2 \ - logger.tensorboard_enabled=true \ - logger.log_dir=$LOG_DIR \ - logger.wandb_enabled=false \ - logger.monitor_gpus=true \ - checkpointing.enabled=true \ - checkpointing.checkpoint_dir=$CKPT_DIR \ - checkpointing.save_period=2 \ - $@ \ - 2>&1 | tee $RUN_LOG - -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS - -# Smoke-level threshold (matches grpo_megatron_generation_async_gym.sh); tighten after CI runs. -uv run tests/check_metrics.py $JSON_METRICS \ - 'median(data["train/gen_kl_error"]) < 1.3' \ - '"3" in data["train/loss"]' \ - '"3" in data["validation/accuracy"]' - -# The save-bound step must defer the engine wake past the checkpoint save. -# `val_period=3` gives us a step 2 that does not wake/sleep cycle the engine before save. -if ! grep -q "Keeping colocated engine asleep for checkpointing" $RUN_LOG; then - echo "FAIL: deferred-wake log line not found (colocated checkpoint path not exercised)" - exit 1 -fi - -if [[ ! -f $CKPT_DIR/step_2/replay_buffer.pt ]]; then - echo "FAIL: replay_buffer.pt not found in step_2 checkpoint" - exit 1 -fi diff --git a/tests/functional/grpo_megatron_generation_colocated_gym.sh b/tests/functional/grpo_megatron_generation_colocated_gym.sh index ad9de24d58a..3326fe2bea3 100755 --- a/tests/functional/grpo_megatron_generation_colocated_gym.sh +++ b/tests/functional/grpo_megatron_generation_colocated_gym.sh @@ -89,7 +89,8 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS -# Smoke-level thresholds, mirroring grpo_megatron_generation_async_gym.sh. +# Lag-0 run: strict engine/trainer token parity on top of the standard gym gates. uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ 'median(data["train/gen_kl_error"]) < 1.3' \ 'data["validation/accuracy"]["10"] > 0.1' diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh b/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh deleted file mode 100644 index 62c8e57f862..00000000000 --- a/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh +++ /dev/null @@ -1,92 +0,0 @@ -#!/bin/bash - -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) -# Mark the current repo as safe, since wandb fetches metadata about the repo -git config --global --add safe.directory $PROJECT_ROOT - -set -eou pipefail - -EXP_NAME=$(basename $0 .sh) -EXP_DIR=$SCRIPT_DIR/$EXP_NAME -LOG_DIR=$EXP_DIR/logs -JSON_METRICS=$EXP_DIR/metrics.json -RUN_LOG=$EXP_DIR/run.log -CKPT_DIR=$EXP_DIR/ckpts -export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} - -rm -rf $EXP_DIR $LOG_DIR -mkdir -p $EXP_DIR $LOG_DIR $CKPT_DIR - -# clean up checkpoint directory on exit -trap "rm -rf $CKPT_DIR" EXIT - -# async colocated, reshard mode: TE TP2 training; inference_optimized TP1 -# generation on a dedicated model, resharded into on every wake. -cd $PROJECT_ROOT -uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ - $PROJECT_ROOT/examples/run_grpo.py \ - --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \ - policy.model_name=Qwen/Qwen3-0.6B \ - grpo.num_prompts_per_step=2 \ - grpo.num_generations_per_prompt=4 \ - policy.train_global_batch_size=4 \ - policy.logprob_batch_size=4 \ - policy.train_micro_batch_size=1 \ - policy.megatron_cfg.tensor_model_parallel_size=2 \ - policy.generation.backend=megatron \ - ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized \ - ++policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ - policy.generation.mcore_generation_config.refit_backend=nccl \ - grpo.async_grpo.enabled=true \ - grpo.async_grpo.max_trajectory_age_steps=1 \ - grpo.async_grpo.in_flight_weight_updates=true \ - loss_fn.use_importance_sampling_correction=true \ - grpo.max_num_steps=3 \ - grpo.val_period=1 \ - grpo.max_val_samples=8 \ - grpo.val_batch_size=8 \ - cluster.gpus_per_node=2 \ - logger.tensorboard_enabled=true \ - logger.log_dir=$LOG_DIR \ - logger.wandb_enabled=false \ - logger.monitor_gpus=true \ - checkpointing.enabled=true \ - checkpointing.checkpoint_dir=$CKPT_DIR \ - checkpointing.save_period=2 \ - $@ \ - 2>&1 | tee $RUN_LOG - -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS - -# Smoke-level threshold (matches grpo_megatron_generation_async_gym.sh); tighten after CI runs. -uv run tests/check_metrics.py $JSON_METRICS \ - 'median(data["train/gen_kl_error"]) < 1.3' \ - '"3" in data["train/loss"]' \ - '"2" in data["validation/accuracy"]' - -# The dedicated inference model must actually be built — guard against this -# leg silently degenerating to the matched-impl (reshardless) path. -if ! grep -q "\[colocated-reshard\] building dedicated inference model" $RUN_LOG; then - echo "FAIL: dedicated-model build log line not found (reshard path not exercised)" - exit 1 -fi - -# The non-save validation (step 1) wakes an already-serving engine; the -# worker must skip it (guards against redundant per-validation resharding). -if ! grep -q "prepare_for_generation: engine already serving, skipping" $RUN_LOG; then - echo "FAIL: idempotent-wake skip log line not found (redundant reshard on validation?)" - exit 1 -fi - -# The save-bound step must defer the engine wake past the checkpoint save. -# With `val_period=1`, the validation always intervenes before the save. -if ! grep -q "Keeping colocated engine asleep for checkpointing" $RUN_LOG; then - echo "FAIL: deferred-wake log line not found (colocated checkpoint path not exercised)" - exit 1 -fi - -if [[ ! -f $CKPT_DIR/step_2/replay_buffer.pt ]]; then - echo "FAIL: replay_buffer.pt not found in step_2 checkpoint" - exit 1 -fi diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard_async_gym.sh b/tests/functional/grpo_megatron_generation_colocated_reshard_async_gym.sh new file mode 100755 index 00000000000..0412b08fe9f --- /dev/null +++ b/tests/functional/grpo_megatron_generation_colocated_reshard_async_gym.sh @@ -0,0 +1,116 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +EXP_NAME=$(basename $0 .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +LOG_DIR=$EXP_DIR/logs +JSON_METRICS=$EXP_DIR/metrics.json +RUN_LOG=$EXP_DIR/run.log +CHECKPOINT_DIR=$EXP_DIR/checkpoints +DATA_DIR=$EXP_DIR/data +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR $CHECKPOINT_DIR $DATA_DIR + +# clean up checkpoint directory on exit +trap "rm -rf $CHECKPOINT_DIR" EXIT + +cd $PROJECT_ROOT + +# Follow nemo-gym instructions here to get this data: +# https://docs.nvidia.com/nemo/gym/0.1.0/tutorials/nemo-rl-grpo/setup.html#training-nemo-rl-grpo-setup +cd 3rdparty/Gym-workspace/Gym + +# We need HF_TOKEN to download the data from huggingface +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi + +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# This trimming of the workplace assistant dataset is necessary b/c with all the tools the first prompt is >4000 tokens +# which will cause vllm to return nothing on the first prompt and crash RL. Since we want to keep this test short to +# smoke test, we trim all but the first tool +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $PROJECT_ROOT/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml \ + policy.model_name=Qwen/Qwen3-0.6B \ + policy.dtensor_cfg.enabled=false \ + policy.megatron_cfg.enabled=true \ + policy.megatron_cfg.tensor_model_parallel_size=2 \ + policy.megatron_cfg.pipeline_model_parallel_size=1 \ + policy.megatron_cfg.expert_model_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=false \ + policy.generation.backend=megatron \ + policy.generation.mcore_generation_config.expose_http_server=true \ + policy.generation.mcore_generation_config.enable_prefix_caching=true \ + policy.max_total_sequence_length=512 \ + policy.generation.max_new_tokens=128 \ + policy.generation.colocated.enabled=true \ + ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized \ + ++policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ + grpo.num_prompts_per_step=4 \ + grpo.num_generations_per_prompt=2 \ + grpo.max_num_steps=10 \ + grpo.val_period=5 \ + grpo.async_grpo.enabled=true \ + grpo.async_grpo.max_trajectory_age_steps=1 \ + grpo.async_grpo.in_flight_weight_updates=true \ + policy.train_global_batch_size=4 \ + policy.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + loss_fn.use_importance_sampling_correction=true \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=true \ + checkpointing.save_period=5 \ + checkpointing.checkpoint_dir=$CHECKPOINT_DIR \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Smoke-level thresholds. Tighten after first successful runs on CI. +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'data["validation/accuracy"]["10"] > 0.1' + +# The save-bound step must defer the engine wake past the checkpoint save +# (save_period=5 with 10 steps gives save-bound steps 5 and 10). +if ! grep -q "Keeping colocated engine asleep for checkpointing" $RUN_LOG; then + echo "FAIL: deferred-wake log line not found (colocated checkpoint path not exercised)" + exit 1 +fi + +# The dedicated inference model must actually be built — guard against this +# leg silently degenerating to the matched-impl (reshardless) path. +if ! grep -q "\[colocated-reshard\] building dedicated inference model" $RUN_LOG; then + echo "FAIL: dedicated-model build log line not found (reshard path not exercised)" + exit 1 +fi diff --git a/tests/functional/grpo_megatron_generation_gym_single_controller.sh b/tests/functional/grpo_megatron_generation_gym_single_controller.sh index e69c2c8017d..f04b16a98a1 100755 --- a/tests/functional/grpo_megatron_generation_gym_single_controller.sh +++ b/tests/functional/grpo_megatron_generation_gym_single_controller.sh @@ -73,7 +73,8 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS -# Observed to be between 0.8-1.3 +# Lag-0 run: strict engine/trainer token parity on top of the standard gym gates uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ 'median(data["train/gen_kl_error"]) < 1.3' \ 'max(data["train/reward"]) > 0' diff --git a/tests/functional/grpo_megatron_generation_async_gym.sh b/tests/functional/grpo_megatron_generation_non_colocated_async_gym.sh similarity index 100% rename from tests/functional/grpo_megatron_generation_async_gym.sh rename to tests/functional/grpo_megatron_generation_non_colocated_async_gym.sh diff --git a/tests/functional/grpo_megatron_generation_non_colocated_gym.sh b/tests/functional/grpo_megatron_generation_non_colocated_gym.sh new file mode 100755 index 00000000000..50a8edf635f --- /dev/null +++ b/tests/functional/grpo_megatron_generation_non_colocated_gym.sh @@ -0,0 +1,99 @@ +#!/bin/bash + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +EXP_NAME=$(basename $0 .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +LOG_DIR=$EXP_DIR/logs +JSON_METRICS=$EXP_DIR/metrics.json +RUN_LOG=$EXP_DIR/run.log +CHECKPOINT_DIR=$EXP_DIR/checkpoints +DATA_DIR=$EXP_DIR/data +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR $CHECKPOINT_DIR $DATA_DIR + +# clean up checkpoint directory on exit +trap "rm -rf $CHECKPOINT_DIR" EXIT + +cd $PROJECT_ROOT + +# Follow nemo-gym instructions here to get this data: +# https://docs.nvidia.com/nemo/gym/0.1.0/tutorials/nemo-rl-grpo/setup.html#training-nemo-rl-grpo-setup +cd 3rdparty/Gym-workspace/Gym + +# We need HF_TOKEN to download the data from huggingface +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi + +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# This trimming of the workplace assistant dataset is necessary b/c with all the tools the first prompt is >4000 tokens +# which will cause vllm to return nothing on the first prompt and crash RL. Since we want to keep this test short to +# smoke test, we trim all but the first tool +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $PROJECT_ROOT/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml \ + policy.model_name=Qwen/Qwen3-0.6B \ + policy.dtensor_cfg.enabled=false \ + policy.megatron_cfg.enabled=true \ + policy.megatron_cfg.tensor_model_parallel_size=1 \ + policy.megatron_cfg.pipeline_model_parallel_size=1 \ + policy.megatron_cfg.expert_model_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=false \ + policy.generation.backend=megatron \ + policy.generation.mcore_generation_config.expose_http_server=true \ + policy.max_total_sequence_length=512 \ + policy.generation.max_new_tokens=128 \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=1 \ + policy.generation.mcore_generation_config.refit_backend=nccl \ + grpo.num_prompts_per_step=4 \ + grpo.num_generations_per_prompt=2 \ + grpo.max_num_steps=10 \ + grpo.val_period=5 \ + policy.train_global_batch_size=4 \ + policy.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=true \ + checkpointing.save_period=5 \ + checkpointing.checkpoint_dir=$CHECKPOINT_DIR \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Lag-0 run: strict engine/trainer token parity on top of the standard gym gates. +uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'data["validation/accuracy"]["10"] > 0.1' diff --git a/tests/functional/grpo_megatron_generation_topp_topk.sh b/tests/functional/grpo_megatron_generation_topp_topk.sh index 7c3ec97ee12..11911693582 100644 --- a/tests/functional/grpo_megatron_generation_topp_topk.sh +++ b/tests/functional/grpo_megatron_generation_topp_topk.sh @@ -47,4 +47,4 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/token_mult_prob_error"]) < 1.06' + 'max(data["train/token_mult_prob_error"]) < 1.10' diff --git a/tests/functional/grpo_megatron_mxfp8_refit_gb200.sh b/tests/functional/grpo_megatron_mxfp8_refit_gb200.sh deleted file mode 100644 index 16c9f7391b4..00000000000 --- a/tests/functional/grpo_megatron_mxfp8_refit_gb200.sh +++ /dev/null @@ -1,117 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -#!/bin/bash -set -euo pipefail - -SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd) -PROJECT_ROOT=$(realpath "$SCRIPT_DIR/../..") -git config --global --add safe.directory "$PROJECT_ROOT" - -EXP_NAME=$(basename "$0" .sh) -EXP_DIR="$SCRIPT_DIR/$EXP_NAME" -LOG_DIR="$EXP_DIR/logs" -JSON_METRICS="$EXP_DIR/metrics.json" -RUN_LOG="$EXP_DIR/run.log" -export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" - -rm -rf "$EXP_DIR" -mkdir -p "$LOG_DIR" - -assert_grep() { - local pattern=$1 - local file=$2 - grep -Eq "$pattern" "$file" || { - echo "[FAIL] expected '$pattern' in $file" - exit 1 - } -} - -cd "$PROJECT_ROOT" -uv run coverage run -a --data-file="$PROJECT_ROOT/tests/.coverage" --source="$PROJECT_ROOT/nemo_rl" \ - "$PROJECT_ROOT/examples/run_grpo.py" \ - --config "$PROJECT_ROOT/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.yaml" \ - grpo.num_prompts_per_step=2 \ - grpo.num_generations_per_prompt=4 \ - grpo.max_num_steps=2 \ - env.math.num_workers=2 \ - loss_fn.reference_policy_kl_penalty=0.0 \ - policy.train_global_batch_size=8 \ - policy.train_micro_batch_size=1 \ - policy.logprob_batch_size=1 \ - policy.max_total_sequence_length=512 \ - policy.make_sequence_length_divisible_by=32 \ - ++policy.megatron_cfg.train_iters=2 \ - policy.megatron_cfg.tensor_model_parallel_size=1 \ - policy.megatron_cfg.expert_model_parallel_size=2 \ - policy.megatron_cfg.sequence_parallel=false \ - policy.megatron_cfg.activation_checkpointing=true \ - policy.megatron_cfg.fp8_cfg.enabled=false \ - policy.megatron_cfg.optimizer.optimizer_cpu_offload=true \ - policy.megatron_cfg.optimizer.optimizer_offload_fraction=1.0 \ - policy.generation.backend=megatron \ - policy.generation.max_new_tokens=64 \ - policy.generation.colocated.enabled=false \ - policy.generation.colocated.resources.gpus_per_node=2 \ - policy.generation.colocated.resources.num_nodes=1 \ - policy.generation.mcore_generation_config.transformer_impl=inference_optimized \ - policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ - policy.generation.mcore_generation_config.expert_model_parallel_size=2 \ - policy.generation.mcore_generation_config.sequence_parallel=false \ - policy.generation.mcore_generation_config.inference_grouped_gemm_backend=torch \ - ++policy.generation.mcore_generation_config.inference_moe_token_dispatcher_type=nvls \ - policy.generation.mcore_generation_config.cuda_graph_impl=local \ - policy.generation.mcore_generation_config.inference_cuda_graph_scope=block \ - policy.generation.mcore_generation_config.num_cuda_graphs=-1 \ - policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=true \ - policy.generation.mcore_generation_config.enable_chunked_prefill=false \ - policy.generation.mcore_generation_config.buffer_size_gb=2 \ - policy.generation.mcore_generation_config.max_model_len=512 \ - policy.generation.mcore_generation_config.max_tokens=512 \ - policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs \ - policy.generation.mcore_generation_config.refit_backend=nccl \ - ++policy.generation.mcore_generation_config.fp8_cfg.enabled=true \ - ++policy.generation.mcore_generation_config.fp8_cfg.fp8=e4m3 \ - ++policy.generation.mcore_generation_config.fp8_cfg.fp8_recipe=mxfp8 \ - ++policy.generation.mcore_generation_config.fp8_cfg.fp8_param=true \ - cluster.gpus_per_node=4 \ - cluster.num_nodes=1 \ - logger.tensorboard_enabled=true \ - logger.log_dir="$LOG_DIR" \ - logger.wandb_enabled=false \ - logger.monitor_gpus=false \ - checkpointing.enabled=false \ - "$@" \ - 2>&1 | tee "$RUN_LOG" - -uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" - -# The setup weight-sync timing proves that the initial refit ran, while the -# per-step transfer timing proves that the post-update refit ran. Raw rollout -# log-probs are compared against BF16 policy recomputation; the bounds allow the -# established MXFP8 quantization delta while rejecting a bad refit. -uv run tests/check_metrics.py "$JSON_METRICS" \ - 'len(data["train/loss"]) == 2' \ - 'len(data["timing/setup/weight_sync_time_s"]) == 1' \ - 'min(data["timing/setup/weight_sync_time_s"]) > 0' \ - 'len(data["timing/train/prepare_for_generation/transfer_and_update_weights"]) == 1' \ - 'min(data["timing/train/prepare_for_generation/transfer_and_update_weights"]) > 0' \ - 'len(data["train/gen_kl_error"]) == 2' \ - 'max(data["train/gen_kl_error"]) < 0.15' \ - 'len(data["train/token_mult_prob_error"]) == 2' \ - 'max(data["train/token_mult_prob_error"]) < 1.5' - -assert_grep 'cuda graph warmup' "$RUN_LOG" - -echo "[PASS] GB200 Nano-v3 BF16-to-MXFP8 Megatron refit functional test" diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index 44e2cf0e5af..622f8d6a965 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -47,3 +47,18 @@ tests/test_suites/vlm/mopd-nemotron-super-omni-120ba12b-4n8g-smoke.v1.sh # nightly budget is at 3922 of its 3928 GPU-hour cap and this run's ~27 # GPU-hours would exceed it. Move to nightly.txt when the budget has room. tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-ready-first-single-controller.sh + +# OOMs on H100 2n8g during the colocated reshard: the colocated training state +# offload is currently too slow to use; tracked by issue #3976. +tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.sh + +# Disabled until NVIDIA/Megatron-LM#6993 is live in this codebase. +tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.sh + +# Megatron-Core/Bridge bump-lane smoke legs. Not disabled: they run in the +# bump-validation lanes via nightly_mcore{,_gb200}.txt rather than in the +# recurring suites above, and are listed here to keep the accounting 1:1. +tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick.sh diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron_generation.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron_generation.sh index b1cc036ae02..fa806860782 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron_generation.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron_generation.sh @@ -39,7 +39,7 @@ if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | ma 'median(data["train/token_mult_prob_error"]) < 1.1' \ 'data["train/token_mult_prob_error"]["500"] < 1.1' \ 'data["train/reward"]["500"] > 0.1' \ - 'mean(data["timing/train/total_step_time"], -6, -1) < 25' + 'median(data["timing/train/total_step_time"]) < 20' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.sh b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.sh index 7806a567472..f729308a5dd 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.sh +++ b/tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.sh @@ -34,12 +34,11 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS # Only run metrics if the target step is reached if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then - # total_step_time observed around ~16, so 17.5 for buffer uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ 'data["train/token_mult_prob_error"]["500"] < 1.1' \ 'data["train/reward"]["500"] > 0.1' \ - 'mean(data["timing/train/total_step_time"], -6, -1) < 17.5' + 'median(data["timing/train/total_step_time"]) < 12' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick.sh new file mode 100755 index 00000000000..ac4042a222c --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick.sh @@ -0,0 +1,67 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=8 +STEPS_PER_RUN=3 +MAX_STEPS=3 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +# Smoke sizing: ~25 min startup + 3 async steps. Deliberately small — this +# leg exists for the mcore-triggered lanes. +NUM_MINUTES=100 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare nemo-gym workplace_assistant dataset (mirrors tests/functional/grpo_async_gym.sh). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +cd 3rdparty/Gym-workspace/Gym +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# Trim to a single tool so first prompt fits. +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +# Run the experiment via the gym entrypoint +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.sh new file mode 100755 index 00000000000..0a08ca2fd3f --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.sh @@ -0,0 +1,71 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=8 +STEPS_PER_RUN=8 +MAX_STEPS=8 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +# ~25 min startup plus 8 async steps; colocated steps carry the engine +# sleep/wake and per-wake reshard on top of training, so keep the async-gym +# sibling's 180 min budget until measured. +NUM_MINUTES=180 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare nemo-gym workplace_assistant dataset (mirrors tests/functional/grpo_async_gym.sh). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +cd 3rdparty/Gym-workspace/Gym +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# Trim to a single tool so first prompt fits. +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +# Run the experiment via the gym entrypoint +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + checkpointing.save_period=4 \ + grpo.val_period=4 \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'max(data["train/reward"]) > 0.0' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick.sh new file mode 100755 index 00000000000..ac4042a222c --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick.sh @@ -0,0 +1,67 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=8 +STEPS_PER_RUN=3 +MAX_STEPS=3 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +# Smoke sizing: ~25 min startup + 3 async steps. Deliberately small — this +# leg exists for the mcore-triggered lanes. +NUM_MINUTES=100 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare nemo-gym workplace_assistant dataset (mirrors tests/functional/grpo_async_gym.sh). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +cd 3rdparty/Gym-workspace/Gym +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# Trim to a single tool so first prompt fits. +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +# Run the experiment via the gym entrypoint +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.sh similarity index 94% rename from tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh rename to tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.sh index 1819c65df23..4bc4be4b55e 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.sh @@ -53,6 +53,7 @@ uv run examples/nemo_gym/run_grpo_nemo_gym.py \ logger.tensorboard_enabled=True \ checkpointing.enabled=True \ checkpointing.checkpoint_dir=$CKPT_DIR \ + checkpointing.save_period=8 \ data.train.data_path=$TRAIN_PATH \ data.validation.data_path=$VALIDATION_PATH \ $@ \ @@ -62,7 +63,8 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/gen_kl_error"]) < 1.3' \ - 'max(data["train/reward"]) > 0.0' + 'max(data["train/reward"]) > 0.0' \ + 'median(data["timing/train/total_step_time"]) < 1000' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.sh deleted file mode 100755 index 9e308ca457e..00000000000 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.sh +++ /dev/null @@ -1,43 +0,0 @@ -#!/bin/bash -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -source $SCRIPT_DIR/common.env - -# ===== BEGIN CONFIG ===== -NUM_NODES=2 -GPUS_PER_NODE=8 -STEPS_PER_RUN=10 -MAX_STEPS=10 -NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=60 -# ===== END CONFIG ===== - -exit_if_max_steps_reached - -# Run the experiment -cd $PROJECT_ROOT -uv run examples/run_grpo.py \ - --config $CONFIG_PATH \ - grpo.max_num_steps=$MAX_STEPS \ - policy.generation.backend=megatron \ - logger.log_dir=$LOG_DIR \ - logger.wandb_enabled=True \ - logger.wandb.project=nemo-rl \ - logger.wandb.name=$EXP_NAME \ - logger.monitor_gpus=True \ - logger.tensorboard_enabled=True \ - checkpointing.enabled=True \ - checkpointing.checkpoint_dir=$CKPT_DIR \ - $@ \ - 2>&1 | tee $RUN_LOG - -# Convert tensorboard logs to json -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS - -# Only run metrics if the target step is reached -if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then - uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/reward"]) > 0.0' - - # Clean up checkpoint directory after successful run to save space. - rm -rf "$CKPT_DIR" -fi diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh deleted file mode 100755 index 2bbe91651a4..00000000000 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh +++ /dev/null @@ -1,37 +0,0 @@ -#!/bin/bash -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -source $SCRIPT_DIR/common.env - -# ===== BEGIN CONFIG ===== -NUM_NODES=4 -GPUS_PER_NODE=4 -SEGMENT_SIZE=2 -STEPS_PER_RUN=10 -MAX_STEPS=10 -NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=60 -# ===== END CONFIG ===== - -exit_if_max_steps_reached - -# Run the experiment -cd $PROJECT_ROOT -uv run examples/run_grpo.py \ - --config $CONFIG_PATH \ - grpo.max_num_steps=$MAX_STEPS \ - logger.log_dir=$LOG_DIR \ - logger.wandb_enabled=True \ - logger.wandb.project=nemo-rl \ - logger.wandb.name=$EXP_NAME \ - logger.monitor_gpus=True \ - logger.tensorboard_enabled=True \ - checkpointing.enabled=False \ - $@ \ - 2>&1 | tee $RUN_LOG - -# Convert tensorboard logs to json -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS - -uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/reward"]) > 0.0' \ - 'median(data["train/gen_kl_error"]) < 1.3' diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick.sh new file mode 100755 index 00000000000..2b056c62c95 --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick.sh @@ -0,0 +1,68 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=4 +SEGMENT_SIZE=2 +STEPS_PER_RUN=3 +MAX_STEPS=3 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +# Smoke sizing: ~25 min startup + 3 async steps. Deliberately small — this +# leg exists for the mcore-triggered lanes. +NUM_MINUTES=100 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare nemo-gym workplace_assistant dataset (mirrors tests/functional/grpo_async_gym.sh). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +cd 3rdparty/Gym-workspace/Gym +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# Trim to a single tool so first prompt fits. +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +# Run the experiment via the gym entrypoint +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.sh new file mode 100755 index 00000000000..94578450254 --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.sh @@ -0,0 +1,72 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=4 +SEGMENT_SIZE=2 +STEPS_PER_RUN=8 +MAX_STEPS=8 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +# ~25 min startup plus 8 async steps; colocated steps carry the engine +# sleep/wake and per-wake reshard on top of training, so keep the async-gym +# sibling's 180 min budget until measured. +NUM_MINUTES=180 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare nemo-gym workplace_assistant dataset (mirrors tests/functional/grpo_async_gym.sh). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +cd 3rdparty/Gym-workspace/Gym +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# Trim to a single tool so first prompt fits. +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +# Run the experiment via the gym entrypoint +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + checkpointing.save_period=4 \ + grpo.val_period=4 \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'max(data["train/reward"]) > 0.0' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick.sh new file mode 100755 index 00000000000..2b056c62c95 --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick.sh @@ -0,0 +1,68 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=4 +SEGMENT_SIZE=2 +STEPS_PER_RUN=3 +MAX_STEPS=3 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +# Smoke sizing: ~25 min startup + 3 async steps. Deliberately small — this +# leg exists for the mcore-triggered lanes. +NUM_MINUTES=100 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare nemo-gym workplace_assistant dataset (mirrors tests/functional/grpo_async_gym.sh). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +cd 3rdparty/Gym-workspace/Gym +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# Trim to a single tool so first prompt fits. +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +# Run the experiment via the gym entrypoint +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.sh new file mode 100755 index 00000000000..12036176a97 --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.sh @@ -0,0 +1,71 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=4 +GPUS_PER_NODE=4 +SEGMENT_SIZE=2 +STEPS_PER_RUN=8 +MAX_STEPS=8 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +# ~25 min startup (30B-MoE load + CUDA-graph warmup + nemo_gym servers) plus 8 +# async steps at ~21 min/step on GB200 4n4g. +NUM_MINUTES=240 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# Prepare nemo-gym workplace_assistant dataset (mirrors tests/functional/grpo_async_gym.sh). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +cd 3rdparty/Gym-workspace/Gym +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# Trim to a single tool so first prompt fits. +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +# Run the experiment via the gym entrypoint +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + checkpointing.save_period=8 \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'max(data["train/reward"]) > 0.0' \ + 'median(data["timing/train/total_step_time"]) < 1500' + +# Clean up checkpoint directory after successful run to save space. +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.sh similarity index 53% rename from tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh rename to tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.sh index 77974b3e33d..9eb7a800c50 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.sh @@ -9,7 +9,8 @@ SEGMENT_SIZE=2 STEPS_PER_RUN=10 MAX_STEPS=10 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=60 +# ~25 min startup (30B load + MXFP8 CUDA-graph warmup) + 10 sync steps; until measured. +NUM_MINUTES=120 # ===== END CONFIG ===== exit_if_max_steps_reached @@ -34,6 +35,20 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS # Only run metrics if the target step is reached if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + # Refit-timing proofs (setup sync + post-update transfers ran) plus the + # MXFP8 quantization-delta band: raw rollout logprobs vs BF16 recompute. uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/reward"]) > 0.0' + 'min(data["timing/setup/weight_sync_time_s"]) > 0' \ + 'min(data["timing/train/prepare_for_generation/transfer_and_update_weights"]) > 0' \ + 'max(data["train/gen_kl_error"]) < 0.15' \ + 'median(data["train/token_mult_prob_error"]) < 1.5' \ + "data['train/token_mult_prob_error']['${MAX_STEPS}'] < 2.0" \ + 'max(data["train/reward"]) > 0.0' \ + 'median(data["timing/train/total_step_time"]) < 48' + + # Generation CUDA graphs must actually engage under MXFP8. + if ! grep -Eq "cuda graph warmup" $RUN_LOG; then + echo "FAIL: 'cuda graph warmup' not found in run log (generation CUDA graphs not engaged)" + exit 1 + fi fi diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.sh deleted file mode 100755 index 2bbe91651a4..00000000000 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.sh +++ /dev/null @@ -1,37 +0,0 @@ -#!/bin/bash -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -source $SCRIPT_DIR/common.env - -# ===== BEGIN CONFIG ===== -NUM_NODES=4 -GPUS_PER_NODE=4 -SEGMENT_SIZE=2 -STEPS_PER_RUN=10 -MAX_STEPS=10 -NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=60 -# ===== END CONFIG ===== - -exit_if_max_steps_reached - -# Run the experiment -cd $PROJECT_ROOT -uv run examples/run_grpo.py \ - --config $CONFIG_PATH \ - grpo.max_num_steps=$MAX_STEPS \ - logger.log_dir=$LOG_DIR \ - logger.wandb_enabled=True \ - logger.wandb.project=nemo-rl \ - logger.wandb.name=$EXP_NAME \ - logger.monitor_gpus=True \ - logger.tensorboard_enabled=True \ - checkpointing.enabled=False \ - $@ \ - 2>&1 | tee $RUN_LOG - -# Convert tensorboard logs to json -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS - -uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/reward"]) > 0.0' \ - 'median(data["train/gen_kl_error"]) < 1.3' diff --git a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.sh similarity index 70% rename from tests/test_suites/llm/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated.sh rename to tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.sh index 83b1eacfacf..5ba045c4ea4 100755 --- a/tests/test_suites/llm/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.sh @@ -3,40 +3,43 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source $SCRIPT_DIR/common.env # ===== BEGIN CONFIG ===== -NUM_NODES=2 -STEPS_PER_RUN=30 -MAX_STEPS=30 +NUM_NODES=1 +GPUS_PER_NODE=4 +# Megatron Inference decodes slower than the vLLM twin (which runs 450 steps +# in 240 minutes); size like the classic megatron_generation nightlies. +STEPS_PER_RUN=50 +MAX_STEPS=50 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up -NUM_MINUTES=120 +NUM_MINUTES=180 # ===== END CONFIG ===== exit_if_max_steps_reached # Run the experiment cd $PROJECT_ROOT -uv run examples/run_grpo.py \ +uv run examples/run_grpo_single_controller.py \ --config $CONFIG_PATH \ grpo.max_num_steps=$MAX_STEPS \ - policy.generation.backend=megatron \ logger.log_dir=$LOG_DIR \ logger.wandb_enabled=True \ logger.wandb.project=nemo-rl \ logger.wandb.name=$EXP_NAME \ logger.monitor_gpus=True \ logger.tensorboard_enabled=True \ - checkpointing.enabled=True \ + checkpointing.enabled=False \ checkpointing.checkpoint_dir=$CKPT_DIR \ $@ \ 2>&1 | tee $RUN_LOG -# Convert tensorboard logs to json uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS # Only run metrics if the target step is reached if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ - 'data["train/token_mult_prob_error"]["30"] < 1.1' + 'data["train/token_mult_prob_error"]["50"] < 1.1' \ + 'max(data["train/reward"]) > 0' \ + 'median(data["timing/train/total_step_time"]) < 30' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-sync.sh b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-sync.sh index 7b6b3f2c124..563365faf28 100755 --- a/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-sync.sh +++ b/tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron_generation-noncolocated-single-controller-sync.sh @@ -34,12 +34,11 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS # Only run metrics if the target step is reached if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then - # Same logprob-health gate as the vLLM twin; no step-time assertion until - # a few runs calibrate Megatron Inference's cadence at this scale. uv run tests/check_metrics.py $JSON_METRICS \ 'median(data["train/token_mult_prob_error"]) < 1.1' \ 'data["train/token_mult_prob_error"]["50"] < 1.1' \ - 'max(data["train/reward"]) > 0' + 'max(data["train/reward"]) > 0' \ + 'median(data["timing/train/total_step_time"]) < 18' # Clean up checkpoint directory after successful run to save space. rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 42f19e86b56..78f679840e7 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -105,7 +105,6 @@ tests/test_suites/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16. # Non-colocated tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated.sh -tests/test_suites/llm/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated.sh tests/test_suites/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.sh # Nemotron Super 49B @@ -121,8 +120,7 @@ tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-fsdp2.sh tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-fsdp2-lora.sh tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-lora.sh tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.sh -tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.sh -tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.sh # Nano-v3.5 tests/test_suites/llm/dapo-nanov3.5-30BA3B-4n8g-automodel.sh diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index a1850e81ace..9ab02ba24ce 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -14,6 +14,8 @@ tests/test_suites/llm/grpo-qwen2.5-7b-instruct-4n4g-fsdp2tp2.v3.sh tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron.sh tests/test_suites/llm/grpo-qwen3-1.7b-1n4g-megatron-eagle3.sh tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron_generation.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.sh # TRT-LLM generation backend tests/test_suites/llm/grpo-qwen3-1.7b-2n4g-fsdp2-trtllm.sh @@ -26,9 +28,6 @@ tests/test_suites/llm/grpo-moonlight-16ba3b-4n4g-megatron.sh tests/test_suites/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.sh tests/test_suites/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.sh -# Nemotron 3 Nano 30B colocated reshard (training/inference layout swap) -tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh - # Functional VLM run tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-dtensor2tp1.v1.sh tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-megatrontp1.v1.sh @@ -50,6 +49,7 @@ tests/test_suites/llm/grpo-nemotron3-super-120BA12B-8n4g-megatron.sh # Single Controller (SC) tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n4g-async-1off-single-controller-streaming2.sh tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron-single-controller-sync.sh +tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatron_generation-noncolocated-single-controller-sync.sh ####### # SFT # @@ -95,9 +95,6 @@ tests/test_suites/llm/distillation-qwen3-32b-to-1.7b-base-1n4g-megatron-tp1pp2cp # Nano3 hybrid MoE/Mamba ModelOpt layer-spec smoke. Keeps # policy.disable_modelopt_layer_spec=false to cover modelopt_mamba_stack_spec. tests/test_suites/llm/distillation-nano3-30ba3b-4n4g-megatron-qa-nvfp4-modelopt-spec.sh -tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh -tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.sh - # SGLang backend tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-fsdp2tp1-sglang.sh tests/test_suites/llm/grpo-qwen2.5-math-1.5b-instruct-1n4g-megatrontp1-sglang.sh diff --git a/tests/test_suites/nightly_mcore.txt b/tests/test_suites/nightly_mcore.txt index d678fff98e5..fef95c0716f 100644 --- a/tests/test_suites/nightly_mcore.txt +++ b/tests/test_suites/nightly_mcore.txt @@ -5,6 +5,8 @@ # Basic megatron tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron.sh tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick.sh # Functional VLM run tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n8g-megatrontp2.v1.sh diff --git a/tests/test_suites/nightly_mcore_gb200.txt b/tests/test_suites/nightly_mcore_gb200.txt index d49864c1199..c3a1248771d 100644 --- a/tests/test_suites/nightly_mcore_gb200.txt +++ b/tests/test_suites/nightly_mcore_gb200.txt @@ -5,6 +5,8 @@ # Basic megatron tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron.sh tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n4g-megatron_generation.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick.sh # Functional VLM run tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-megatrontp1.v1.sh diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh index bc36a391470..0c456ff5f20 100755 --- a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh @@ -9,6 +9,7 @@ source $SCRIPT_DIR/common.env # ===== BEGIN CONFIG ===== NUM_NODES=8 GPUS_PER_NODE=4 +SEGMENT_SIZE=2 STEPS_PER_RUN=50 MAX_STEPS=50 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))