diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.yaml new file mode 100644 index 00000000000..2c61d2b4c6e --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.yaml @@ -0,0 +1,38 @@ +defaults: ../../grpo_math_1B.yaml +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 8 +checkpointing: + checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp +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: 8192 + make_sequence_length_divisible_by: 8 + dtensor_cfg: + enabled: false + megatron_cfg: + enabled: true + bias_activation_fusion: false + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 2 + context_parallel_size: 2 + expert_model_parallel_size: 8 + sequence_parallel: true + generation: + vllm_cfg: + tensor_parallel_size: 4 + gpu_memory_utilization: 0.7 +logger: + wandb_enabled: true + tensorboard_enabled: true + wandb: + project: nemo-rl + name: grpo-nanov3-30BA3B-2n8g-megatron-pack-cp +cluster: + gpus_per_node: 8 + num_nodes: 2 diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index 3e1e2dfec0b..5dee00f678e 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -509,6 +509,8 @@ def _pack_sequences_for_megatron( if cu_seqlens_padded is None: cu_seqlens_padded = cu_seqlens.clone() + # total_tokens is required for PackedSeqParams.__post_init__ to build + # seq_idx, which Mamba uses to reset SSM state at sample boundaries. packed_seq_params = PackedSeqParams( cu_seqlens_q=cu_seqlens_padded, cu_seqlens_kv=cu_seqlens_padded, @@ -517,6 +519,7 @@ def _pack_sequences_for_megatron( max_seqlen_q=int(max_seqlen), max_seqlen_kv=int(max_seqlen), qkv_format="thd", + total_tokens=packed_input_ids.shape[1], ) return ( diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.sh new file mode 100755 index 00000000000..c1ea21693b2 --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.sh @@ -0,0 +1,44 @@ +#!/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 +NUM_MINUTES=30 +# ===== 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=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 \ + 'median(data["train/token_mult_prob_error"]) < 1.1' \ + 'mean(data["train/gen_kl_error"]) < 0.02' \ + 'max(data["train/reward"]) > 0.05' + + # Clean up checkpoint directory after successful run to save space. + rm -rf "$CKPT_DIR" +fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 0ceb4b13fdd..7917083ff71 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -84,6 +84,7 @@ tests/test_suites/llm/grpo-nano-v2-12b-2n8g-fsdp2tp1.sh 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 # Lora tests/test_suites/llm/grpo-qwen3-8B-base-1n8g-fsdp2-lora.sh diff --git a/tests/unit/models/megatron/megatron_data_actors.py b/tests/unit/models/megatron/megatron_data_actors.py index 0631b656392..4dd83886d94 100644 --- a/tests/unit/models/megatron/megatron_data_actors.py +++ b/tests/unit/models/megatron/megatron_data_actors.py @@ -497,6 +497,25 @@ def _test_context_parallel( "error": f"CP wrong qkv_format: expected 'thd', got {packed_seq_params.qkv_format}", } + # Mamba SSM state reset relies on packed_seq_params.seq_idx, which + # __post_init__ only builds when total_tokens is set. + if packed_seq_params.total_tokens != expected_total_tokens: + return { + "success": False, + "error": f"CP packed_seq_params.total_tokens mismatch: expected {expected_total_tokens}, got {packed_seq_params.total_tokens}", + } + if packed_seq_params.seq_idx is None: + return { + "success": False, + "error": "CP packed_seq_params.seq_idx is None", + } + expected_seq_idx_len = int(cu_seqlens_padded[-1].item()) + if packed_seq_params.seq_idx.shape != (1, expected_seq_idx_len): + return { + "success": False, + "error": f"CP packed_seq_params.seq_idx shape mismatch: expected (1, {expected_seq_idx_len}), got {tuple(packed_seq_params.seq_idx.shape)}", + } + # Test 2: CP packing with full sequence padding pad_full_seq_to = (batch_size * seq_len) + 8 # Add some padding ( @@ -532,6 +551,22 @@ def _test_context_parallel( "error": f"CP full padding cu_seqlens_padded mismatch: expected {expected_cu_seqlens_padded_full}, got {cu_seqlens_padded_full}", } + if packed_seq_params_full.total_tokens != expected_tokens_per_rank_full: + return { + "success": False, + "error": f"CP (full pad) packed_seq_params.total_tokens mismatch: expected {expected_tokens_per_rank_full}, got {packed_seq_params_full.total_tokens}", + } + if packed_seq_params_full.seq_idx is None: + return { + "success": False, + "error": "CP (full pad) packed_seq_params.seq_idx is None", + } + if packed_seq_params_full.seq_idx.shape != (1, pad_full_seq_to): + return { + "success": False, + "error": f"CP (full pad) packed_seq_params.seq_idx shape mismatch: expected (1, {pad_full_seq_to}), got {tuple(packed_seq_params_full.seq_idx.shape)}", + } + correct_ids_0 = torch.tensor( [0, 1, 2, 3, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 0, 0, 0, 0, 0, 0], device="cuda",