Training Muse-Spark-Open-816B-A42B on H200

Author

ATTUNE (cite)

GPUs

128 x H200 141 GB (AWS p5en.48xlarge, 8 per node)

Layout

TP 4, PP 16 (first and last stage 3 layers, the 14 middle stages 4 each), CP 1, EP 8, ETP 1, DP 2, sequence parallel; one 8,192-token row per microbatch (THD, max 8,192 tokens/GPU); full recompute (uniform, 1 layer); distributed optimizer; CPU optimizer offload of 30% with overlapped copies and the precision-aware optimizer; alltoall dispatcher

Tokens/s per GPU

333.12

$ / 1M tokens

$6.60

MFU

9.0% (N = 44,432,883,712 active parameters per token)

Step-time CV

0.82%

Convergence check

correlation 0.99574, final loss difference 0.0003, 0 pointwise failures over 20 steps

Component test

CPU optimizer offload with the precision-aware optimizer (fractions 1.0 and 0.5, with and without overlapped copies) against an FP64 reference: H200, 2026-10-16. Shape: one Muse Spark latent MoE layer: hidden 4096 (the latent width), top-8, expert FFN 4096, 128 experts so the FP64 reference fits, EP 8, ETP 1, 4,096 tokens per rank; the Muse Spark router (sigmoid scores, FP32 weight, frozen expert bias). worst tensor (dense update) relative L2 error 7.96e-7 against a bound of 1.0e-3 after one step, and 7.49e-5 against 1.1e-3 after three steps at the recipe's hyperparameters (bound 2 x the GPU optimizer's error + 1e-3); the precision-aware optimizer without offload also passes, the TE remainder self-test passes and the three negative controls fail run on this recipe's training image; fraction 0.3 moves whole parameter shards to the CPU until they hold 30% of the elements, as fraction 0.5 does at half

Measured

2026-10-17

GPU-hours

38.61

Rows

muse-spark-tulu3-fixed8192-v1 (rows_sha256 4003478c605da1a2a351a0ba38adba24bcef54355394cedf4a1d425ec049128f)

export CUDA_DEVICE_MAX_CONNECTIONS=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
export TORCHINDUCTOR_COMPILE_THREADS=1
export OMP_NUM_THREADS=12
export EP_SUPPRESS_NCCL_CHECK=1
export TILELANG_CACHE_DIR="$KERNEL_CACHE/tilelang"
export TRITON_CACHE_DIR="$KERNEL_CACHE/triton"
export TORCHINDUCTOR_CACHE_DIR="$KERNEL_CACHE/inductor"
export FI_PROVIDER=efa
export FI_EFA_USE_DEVICE_RDMA=1

MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py muse-spark-full)"
read -ra MODEL_ARGS <<< "$MODEL_ARGS_LINE"
python3 train.py \
  --actor-num-nodes 16 \
  --actor-num-gpus-per-node 8 \
  --num-gpus-per-node 8 \
  "${MODEL_ARGS[@]}" \
  --hf-checkpoint "$HF_CKPT" \
  --ref-load "$TORCH_DIST" \
  --rollout-function-path pretokenized_sft_rollout.generate_rollout \
  --prompt-data "$SFT_DATA" \
  --input-key messages \
  --metadata-key metadata \
  --num-rollout 20 \
  --rollout-batch-size 128 \
  --global-batch-size 128 \
  --loss-type sft_loss \
  --calculate-per-token-loss \
  --disable-compute-advantages-and-returns \
  --debug-train-only \
  --debug-skip-weight-update \
  --tensor-model-parallel-size 4 \
  --sequence-parallel \
  --pipeline-model-parallel-size 16 \
  --decoder-first-pipeline-num-layers 3 \
  --decoder-last-pipeline-num-layers 3 \
  --context-parallel-size 1 \
  --expert-model-parallel-size 8 \
  --expert-tensor-parallel-size 1 \
  --qkv-format thd \
  --seq-length 8192 \
  --attention-backend flash \
  --use-dynamic-batch-size \
  --max-tokens-per-gpu 8192 \
  --recompute-granularity full \
  --recompute-method uniform \
  --recompute-num-layers 1 \
  --optimizer adam --lr 1e-5 --lr-decay-style cosine --min-lr 1e-6 \
  --lr-warmup-fraction 0.1 --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \
  --use-distributed-optimizer \
  --use-precision-aware-optimizer \
  --optimizer-cpu-offload \
  --optimizer-offload-fraction 0.3 \
  --overlap-cpu-optimizer-d2h-h2d \
  --distributed-timeout-minutes 60

# recipes/muse-spark-open-816b-a42b/h200/attune-tp4-pp16-ep8.yaml
# An OpenRecipes recipe; the fields are documented in schema/recipe.schema.yaml.
schema_version: 1
model: Muse-Spark-Open-816B-A42B
gpu:
  type: H200
  memory_gb: 141
  count: 128
  per_node: 8
  instance: AWS p5en.48xlarge
author:
  name: ATTUNE
  url: https://arxiv.org/abs/2610.17528
status: verified
stack:
  framework: 'miles commit bfdcff0e (radixark/miles PR #4111, the commit of the published radixark/miles:muse-spark-open
    image), no local changes'
  megatron_lm: radixark/Megatron-LM commit 6530834a (branch muse-spark-open; megatron-core 0.19.0), no local changes
  transformer_engine: 2.17.0 (PyTorch extension rebuilt against the installed PyTorch)
  pytorch: 2.14.1 (CUDA 13.0 build)
  cuda: '13.0'
  nccl: 2.30.7 (CUDA 13.3 build; printed by the measured run)
  flash_attn: 2.7.4.post1 (rebuilt against the installed PyTorch)
  deepep: sgl-deep-ep 0.1.2.post1 (installed; the alltoall dispatcher does not use it)
  python: 3.12.3
  network: 8 GPUs per node, EFA 1.49 (aws-ofi-nccl 1.20), 3.2 Tb/s RDMA per node
  source: 'version check of the measured run''s training image (same digest) at its build, 2026-10-16: miles and
    Megatron-LM git commits, installed TE, flash-attn, PyTorch and megatron-core; package listing and git status of the
    same image, 2026-10-18: Python, DeepEP, clean miles and Megatron-LM checkouts; NCCL from the measured run''s log.
    The image is radixark/miles:muse-spark-open (amd64 digest
    sha256:2b46ca326003e347c600145bd60935b42af8032ed6f0a75e7fdc7bf5a8700b15) with EFA 1.49 added, TE''s PyTorch
    extension and flash-attn 2 rebuilt against its PyTorch, and Ubuntu security updates'
workload:
  row_tokens: 8192
  rows: SFT conversations
  global_batch: 128
  precision: bf16
  dataset:
    name: Tulu 3 SFT mixture
    source: allenai/tulu-3-sft-mixture
    revision: b14afda60f1bbebe55d5d2fa1e4df5042f97f8be
    license: ODC-By 1.0
    rows_sha256: 4003478c605da1a2a351a0ba38adba24bcef54355394cedf4a1d425ec049128f
    build:
      name: muse-spark-tulu3-fixed8192-v1
      tokenizer: Muse-Spark-Open-816B-A42B (revision 0fb02478b02565b2ce7fc57e6051c91172774ac9; tokenizer.json sha256
        ac2f41ef2444dec82fdc90e334f80543f298535005fe78aec7e8311ef86f84d0; chat_template.jinja sha256
        2341a2d3b1f1b0ac92b6237d46e2fff495109bd3305cacdea54726f5e98411db)
      loss_mask: muse_spark
      rows: 6400
      row_tokens: 8192
      padding_tokens: 0
      seed: 42
      packing: source conversations in a seeded global permutation, concatenated with an EOS token between them and cut
        into consecutive 8,192-token rows; rows without supervised targets are skipped
      template_kwargs: current_date 2026-10-01, reasoning_effort xhigh
      supervised_token_ratio: 0.56
      parquet_sha256: 19d690622257686f709057f8e665126b8847a18eb6f303e6e9d0d77bdeb1f576
      parquet_bytes: 68772137
      builder: data/generate_fixed_sft.py with --loss-mask-type muse_spark (muse_spark_mask.py; data/README.md)
layout:
  label: TP4 PP16 EP8 DP2
  summary: 'TP 4, PP 16 (first and last stage 3 layers, the 14 middle stages 4 each), CP 1, EP 8, ETP 1, DP 2, sequence
    parallel; one 8,192-token row per microbatch (THD, max 8,192 tokens/GPU); full recompute (uniform, 1 layer);
    distributed optimizer; CPU optimizer offload of 30% with overlapped copies and the precision-aware optimizer;
    alltoall dispatcher'
  highlights: max 8,192 tokens/GPU, full recompute, CPU optimizer offload 0.3
  tp: 4
  pp: 16
  vpp: 1
  cp: 1
  ep: 8
  etp: 1
  dp: 2
  sequence_parallel: true
  layer_split:
    first: 3
    last: 3
  max_tokens_per_gpu: 8192
  rows_per_microbatch: 1
  recompute:
    granularity: full
    method: uniform
    num_layers: 1
  offload:
    optimizer_fraction: 0.3
    overlap_d2h_h2d: true
  dispatcher: alltoall
result:
  tokens_per_s_per_gpu: 333.12
  usd_per_million_tokens: 6.60
  step_time_cv_percent: 0.82
  mfu: 0.0898
  mfu_note: null
  active_params:
    count: 44432883712
    source: https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B
    revision: 0fb02478b02565b2ce7fc57e6051c91172774ac9
  measured: 2026-10-17
  steps_per_run: 20
  window:
  - 5
  - 19
  window_tokens: 15728640
  runs: 1
  profiler: false
  gpu_hours: 38.61
  gpu_hours_note: GPUs x (scheduler admission to finish), including startup; the same run is the candidate arm of the
    convergence check
correctness:
  result: pass
  summary: passed; 1 component test passed on this image
  reference: 'miles'' validated 64-GPU H200 layout for Muse Spark Open (radixark/miles PR #4111): TP 8 with sequence
    parallel, PP 8 (first and last stage 7 layers), EP 8, full recompute, here on 128 GPUs (DP 2) with Adam on the CPU
    (`--optimizer-cpu-offload --optimizer-offload-fraction 1.0 --overlap-cpu-optimizer-d2h-h2d
    --use-precision-aware-optimizer`) in place of the PR''s disk-streamed optimizer state, which needs node-local NVMe
    that the measured nodes did not expose to the job. Model arguments, rows and schedule are the same in both recipes.'
  configuration: pass
  component_tests:
  - test: CPU optimizer offload with the precision-aware optimizer (fractions 1.0 and 0.5, with and without overlapped
      copies) against an FP64 reference
    shape: 'one Muse Spark latent MoE layer: hidden 4096 (the latent width), top-8, expert FFN 4096, 128 experts so the
      FP64 reference fits, EP 8, ETP 1, 4,096 tokens per rank; the Muse Spark router (sigmoid scores, FP32 weight,
      frozen expert bias)'
    gpu: H200
    date: 2026-10-16
    result: pass
    detail: worst tensor (dense update) relative L2 error 7.96e-7 against a bound of 1.0e-3 after one step, and 7.49e-5
      against 1.1e-3 after three steps at the recipe's hyperparameters (bound 2 x the GPU optimizer's error + 1e-3); the
      precision-aware optimizer without offload also passes, the TE remainder self-test passes and the three negative
      controls fail
    note: run on this recipe's training image; fraction 0.3 moves whole parameter shards to the CPU until they hold 30%
      of the elements, as fraction 0.5 does at half
  convergence:
    steps: 20
    result: pass
    correlation: 0.99574
    final_loss_difference: 0.0003
    pointwise_failures: 0
    note: 'candidate arm: the measured run; reference arm: a 20-step run of the reference recipe on the same image and
      rows; the parity probe off in both'
  rules: argument classes of ATTUNE's skill/scripts/correctness/rules/numerics_list.yaml (amazon-science/attune f2cb450),
    with miles' per-run run_uuid, newer than that list, ignored; 20-step check by Megatron-Bridge validate_convergence
    with its pipeline thresholds
env:
  CUDA_DEVICE_MAX_CONNECTIONS: '1'
  PYTORCH_CUDA_ALLOC_CONF: expandable_segments:False
  NVTE_ALLOW_NONDETERMINISTIC_ALGO: '1'
  TORCHINDUCTOR_COMPILE_THREADS: '1'
  OMP_NUM_THREADS: '12'
  EP_SUPPRESS_NCCL_CHECK: '1'
  TILELANG_CACHE_DIR: $KERNEL_CACHE/tilelang
  TRITON_CACHE_DIR: $KERNEL_CACHE/triton
  TORCHINDUCTOR_CACHE_DIR: $KERNEL_CACHE/inductor
  FI_PROVIDER: efa
  FI_EFA_USE_DEVICE_RDMA: '1'
patches: []
run: |
  export CUDA_DEVICE_MAX_CONNECTIONS=1
  export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
  export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
  export TORCHINDUCTOR_COMPILE_THREADS=1
  export OMP_NUM_THREADS=12
  export EP_SUPPRESS_NCCL_CHECK=1
  export TILELANG_CACHE_DIR="$KERNEL_CACHE/tilelang"
  export TRITON_CACHE_DIR="$KERNEL_CACHE/triton"
  export TORCHINDUCTOR_CACHE_DIR="$KERNEL_CACHE/inductor"
  export FI_PROVIDER=efa
  export FI_EFA_USE_DEVICE_RDMA=1

  MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py muse-spark-full)"
  read -ra MODEL_ARGS <<< "$MODEL_ARGS_LINE"
  python3 train.py \
    --actor-num-nodes 16 \
    --actor-num-gpus-per-node 8 \
    --num-gpus-per-node 8 \
    "${MODEL_ARGS[@]}" \
    --hf-checkpoint "$HF_CKPT" \
    --ref-load "$TORCH_DIST" \
    --rollout-function-path pretokenized_sft_rollout.generate_rollout \
    --prompt-data "$SFT_DATA" \
    --input-key messages \
    --metadata-key metadata \
    --num-rollout 20 \
    --rollout-batch-size 128 \
    --global-batch-size 128 \
    --loss-type sft_loss \
    --calculate-per-token-loss \
    --disable-compute-advantages-and-returns \
    --debug-train-only \
    --debug-skip-weight-update \
    --tensor-model-parallel-size 4 \
    --sequence-parallel \
    --pipeline-model-parallel-size 16 \
    --decoder-first-pipeline-num-layers 3 \
    --decoder-last-pipeline-num-layers 3 \
    --context-parallel-size 1 \
    --expert-model-parallel-size 8 \
    --expert-tensor-parallel-size 1 \
    --qkv-format thd \
    --seq-length 8192 \
    --attention-backend flash \
    --use-dynamic-batch-size \
    --max-tokens-per-gpu 8192 \
    --recompute-granularity full \
    --recompute-method uniform \
    --recompute-num-layers 1 \
    --optimizer adam --lr 1e-5 --lr-decay-style cosine --min-lr 1e-6 \
    --lr-warmup-fraction 0.1 --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \
    --use-distributed-optimizer \
    --use-precision-aware-optimizer \
    --optimizer-cpu-offload \
    --optimizer-offload-fraction 0.3 \
    --overlap-cpu-optimizer-d2h-h2d \
    --distributed-timeout-minutes 60
notes: |
  - Starting point: miles' Muse Spark launcher for 64 H200, `scripts/run_muse_spark_64xh200.py`
    (https://github.com/radixark/miles/pull/4111; model page https://miles.radixark.com/docs/models/muse-spark/muse-spark-open).
    miles publishes it as a GRPO recipe with step times and no tokens/s figure. Ported to this workload with its layout
    (TP 8, PP 8 with 7 layers on the first and last stages, EP 8) on 128 H200, with Adam on the CPU in place of its
    NVMe-streamed optimizer state, it measured 224.4 tokens/s/GPU; ATTUNE tuned it to 333.1 (+48%).
  - `OMP_NUM_THREADS=12` gives each training process 12 threads for its part of the CPU Adam step, the node's 96
    physical cores over 8 ranks; Ray otherwise starts each training actor with one thread.
  - Keep the full recompute: a tuning run at PP 16 with recompute off ran out of memory at step 0 with the whole
    optimizer on the CPU, and this recipe keeps 70% of the optimizer state on the GPU.