High token_mult_prob_error when running GRPO on GB200 #1330

grpo:
  num_prompts_per_step: 64
  num_generations_per_prompt: 32
  max_rollout_turns: 1
  max_num_steps: 1000
  normalize_rewards: true
  use_leave_one_out_baseline: true
  val_period: 5
  val_at_start: true
  overlong_filtering: false
  max_val_samples: 128
  val_batch_size: 128
  seed: 42
loss_fn:
  reference_policy_kl_penalty: 0.01
  ratio_clip_min: 0.2
  ratio_clip_max: 0.2
  ratio_clip_c: null
  use_on_policy_kl_approximation: false
  use_importance_sampling_correction: false
  token_level_loss: true
checkpointing:
  enabled: true
  checkpoint_dir: results/setme
  metric_name: val_reward
  higher_is_better: true
  keep_top_k: 100
  save_period: 5
  checkpoint_must_save_by: null
policy:
  model_name: Qwen/Qwen3-32B
  tokenizer:
    name: ${policy.model_name}
    # Custom template that keeps thinking tags
    chat_template: "{%- for message in messages %}{%- if message.role == 'system' %}\n<|im_start|>system\n{{- message.content }}<|im_end|>\n{%- elif message.role == 'user' %}\n<|im_start|>user\n{{- message.content }}<|im_end|><|im_start|>assistant{%- elif message.role == 'assistant' %}\n{{- message.content }}<|im_end|>{%- endif %}{%- endfor %}"
  train_global_batch_size: 128
  train_micro_batch_size: 1
  generation_batch_size: 16
  logprob_batch_size: 2
  max_total_sequence_length: 16384
  precision: bfloat16
  fsdp_offload_enabled: false
  activation_checkpointing_enabled: false
  dtensor_cfg:
    enabled: true
    cpu_offload: false
    sequence_parallel: true
    activation_checkpointing: true
    tensor_parallel_size: 8
    context_parallel_size: 1
    custom_parallel_plan: null
  dynamic_batching:
    enabled: true  # true (changed)
    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
    sequence_length_round: 64
  sequence_packing:
    enabled: false
    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
    algorithm: "modified_first_fit_decreasing"
    sequence_length_round: 64
  make_sequence_length_divisible_by: ${policy.dtensor_cfg.tensor_parallel_size}
  max_grad_norm: 1
  optimizer:
    name: torch.optim.AdamW
    kwargs:
      lr: 3e-07
      weight_decay: 0.01
      betas:
        - 0.9
        - 0.999
      eps: 1e-08
      foreach: false
      fused: false
  scheduler:
    - name: torch.optim.lr_scheduler.LinearLR
      kwargs:
        start_factor: 0.1
        end_factor: 1
        total_iters: 13
    - name: torch.optim.lr_scheduler.ConstantLR
      kwargs:
        factor: 1
        total_iters: 10000000000
    - milestones:
        - 13
  generation:
    backend: vllm
    max_new_tokens: ${policy.max_total_sequence_length}
    temperature: 1
    top_p: 1
    top_k: null
    stop_token_ids: null
    stop_strings: null
    vllm_cfg:
      async_engine: false
      precision: ${policy.precision}
      tensor_parallel_size: 8
      pipeline_parallel_size: 1
      enable_expert_parallel: false
      gpu_memory_utilization: 0.4
      max_model_len: ${policy.max_total_sequence_length}
      enforce_eager: false
    colocated:
      enabled: true
      resources:
        gpus_per_node: null
        num_nodes: null

FROM gitlab-master.nvidia.com/dl/joc/nemo-ci/main/rl:pipe.33623463

# Make local or remote NeMo RL sources available at /opt/nemo-rl
COPY --from=nemo-rl . /opt/nemo-rl

# Install uv and python
ARG UV_VERSION=0.7.2
ARG PYTHON_VERSION=3.12
ENV PATH="/root/.local/bin:$PATH"
RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
    uv python install ${PYTHON_VERSION} && \
    apt-get update && \
    apt-get install -y python${PYTHON_VERSION}-dev && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*