Gptoss 20b swa functional test

gpt_builders.py

@@ -29,6 +29,14 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
             config = core_transformer_config_from_yaml(args, "language_model")
         else:
             config = core_transformer_config_from_args(args)
+    if args.position_embedding_type == 'yarn':
+        config.yarn_rotary_scaling_factor = args.rotary_scaling_factor
+        config.yarn_original_max_position_embeddings = 4096
+        config.yarn_beta_fast = 32.0
+        config.yarn_beta_slow = 1.0
+        config.yarn_mscale = args.mscale
+        config.yarn_mscale_all_dim = args.mscale_all_dim
+        config.yarn_correction_range_round_to_int = False
     if args.spec is not None:
         transformer_layer_spec = import_module(args.spec)
     else:

tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa/golden_values_dev_dgx_gb200.json

@@ -0,0 +1,82 @@
+{
+ "0": {
+  "input_prompt": "The capital of France is",
+  "generated_text": "-13  ( inter- patternEX: ?/ 0\n\n equivalent,",
+  "generated_tokens": [
+   12,
+   1311,
+   220,
+   350,
+   993,
+   12,
+   8302,
+   3922,
+   25,
+   1423,
+   14,
+   220,
+   15,
+   279,
+   23458,
+   11
+  ],
+  "latency": 3.226062774658203,
+  "ttft": 0.20889067649841309,
+  "cuda_graph_request_count_map": null,
+  "step_count": 16,
+  "top_n_logprobs": null,
+  "prompt_top_n_logprobs": null,
+  "prompt_logprobs": [
+   -17.427139282226562,
+   -9.624153137207031,
+   -13.227917671203613,
+   -12.510149002075195
+  ],
+  "generated_logprobs": [
+   -2.727036237716675,
+   -3.0633504390716553,
+   -1.933884859085083,
+   -3.0503389835357666,
+   -2.1997787952423096,
+   -2.5635645389556885,
+   -3.5620317459106445,
+   -2.0540547370910645,
+   -2.0354530811309814,
+   -2.1969399452209473,
+   -1.69447922706604,
+   -1.9973949193954468,
+   -0.8427522778511047,
+   -0.7901788949966431,
+   -2.986577272415161,
+   -2.205671787261963
+  ],
+  "logprobs": [
+   -17.427139282226562,
+   -9.624153137207031,
+   -13.227917671203613,
+   -12.510149002075195,
+   -2.727036237716675,
+   -3.0633504390716553,
+   -1.933884859085083,
+   -3.0503389835357666,
+   -2.1997787952423096,
+   -2.5635645389556885,
+   -3.5620317459106445,
+   -2.0540547370910645,
+   -2.0354530811309814,
+   -2.1969399452209473,
+   -1.69447922706604,
+   -1.9973949193954468,
+   -0.8427522778511047,
+   -0.7901788949966431,
+   -2.986577272415161,
+   -2.205671787261963
+  ]
+ },
+ "throughput": [
+  0.6947979243712443,
+  4.946439130054707
+ ],
+ "mem-max-allocated-bytes": 32378457088,
+ "lifetime_prefill_token_count": 5
+}

tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa/golden_values_dev_dgx_h100.json

@@ -0,0 +1,82 @@
+{
+ "0": {
+  "input_prompt": "The capital of France is",
+  "generated_text": "-13 \n\nUnfortunately 0 up!  0 0- ",
+  "generated_tokens": [
+   12,
+   1311,
+   220,
+   279,
+   51832,
+   220,
+   15,
+   869,
+   0,
+   220,
+   220,
+   15,
+   220,
+   15,
+   12,
+   220
+  ],
+  "latency": 2.3446803092956543,
+  "ttft": 0.21960043907165527,
+  "cuda_graph_request_count_map": null,
+  "step_count": 16,
+  "top_n_logprobs": null,
+  "prompt_top_n_logprobs": null,
+  "prompt_logprobs": [
+   -17.367233276367188,
+   -9.547689437866211,
+   -13.360268592834473,
+   -12.42806339263916
+  ],
+  "generated_logprobs": [
+   -2.7885403633117676,
+   -2.9927821159362793,
+   -1.9823970794677734,
+   -2.99981427192688,
+   -2.5622572898864746,
+   -1.6538726091384888,
+   -1.7417904138565063,
+   -3.610473155975342,
+   -2.025908946990967,
+   -2.3121378421783447,
+   -1.4078569412231445,
+   -0.7797510027885437,
+   -0.8604459762573242,
+   -0.8619584441184998,
+   -1.153270959854126,
+   -0.7719088196754456
+  ],
+  "logprobs": [
+   -17.367233276367188,
+   -9.547689437866211,
+   -13.360268592834473,
+   -12.42806339263916,
+   -2.7885403633117676,
+   -2.9927821159362793,
+   -1.9823970794677734,
+   -2.99981427192688,
+   -2.5622572898864746,
+   -1.6538726091384888,
+   -1.7417904138565063,
+   -3.610473155975342,
+   -2.025908946990967,
+   -2.3121378421783447,
+   -1.4078569412231445,
+   -0.7797510027885437,
+   -0.8604459762573242,
+   -0.8619584441184998,
+   -1.153270959854126,
+   -0.7719088196754456
+  ]
+ },
+ "throughput": [
+  0.9248038506887934,
+  6.811116750769296
+ ],
+ "mem-max-allocated-bytes": 32380357632,
+ "lifetime_prefill_token_count": 5
+}

tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa/model_config.yaml

@@ -0,0 +1,98 @@
+# Inference functional test: GPT-OSS-20B with sliding-window + sink attention (SWA).
+
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: :4096:8
+  HF_HOME: ${DATA_PATH}/hf_home
+
+TEST_TYPE: frozen-start
+MODE: inference
+
+MODEL_ARGS:
+  --use-mcore-models: true
+  --transformer-impl: transformer_engine
+  --distributed-backend: nccl
+
+  # Tokenizer & checkpoint
+  --tokenizer-type: HuggingFaceTokenizer
+  --tokenizer-model: unsloth/gpt-oss-20b-BF16
+  --load: ${CHECKPOINT_LOAD_PATH}/model/openai_gpt-oss-20b/v1
+  --auto-detect-ckpt-format: true
+  --ckpt-format: torch_dist
+  --no-load-optim: true
+  --no-use-tokenizer-model-from-checkpoint-args: true
+  --dist-ckpt-strictness: log_unexpected
+  --inference-ckpt-non-strict: true
+
+  # Parallelism — must match converted checkpoint (TP2 * PP2 * EP2 = 8 GPUs)
+  --tensor-model-parallel-size: 2
+  --pipeline-model-parallel-size: 2
+  --expert-model-parallel-size: 2
+  --expert-tensor-parallel-size: 1
+  --moe-token-dispatcher-type: alltoall
+  --moe-grouped-gemm: true
+
+  # GPT-OSS-20B architecture (matches converted checkpoint)
+  --num-layers: 24
+  --hidden-size: 2880
+  --ffn-hidden-size: 2880
+  --num-attention-heads: 64
+  --group-query-attention: true
+  --num-query-groups: 8
+  --kv-channels: 64
+  --num-experts: 32
+  --moe-ffn-hidden-size: 2880
+  --moe-router-topk: 4
+  --moe-router-dtype: fp32
+  --moe-router-score-function: softmax
+  --moe-router-load-balancing-type: aux_loss
+  --moe-aux-loss-coeff: 0.0
+  --untie-embeddings-and-output-weights: true
+  --disable-bias-linear: true
+  --normalization: RMSNorm
+  --position-embedding-type: yarn
+  --rotary-base: 150000
+  --rotary-percent: 1.0
+  --rotary-scaling-factor: 32.0
+  --quick-geglu: true
+  --glu-linear-offset: 1.0
+  --activation-func-clamp-value: 7.0
+  --softmax-type: learnable
+  --window-size: 127,0
+  --window-attn-skip-freq: 2
+  --padded-vocab-size: 201088
+  --make-vocab-size-divisible-by: 128
+  --seq-length: 4096
+  --max-position-embeddings: 40960
+  --no-rope-fusion: true
+  --no-masked-softmax-fusion: true
+
+  --bf16: true
+  --attention-backend: flash
+  --deterministic-mode: true
+  --micro-batch-size: 1
+
+  # Dynamic inference engine
+  --max-tokens-to-oom: 3600000
+  --inference-max-seq-length: 4096
+  --inference-dynamic-batching-buffer-size-gb: 20
+  --incoming-requests-per-step: 4
+  --inference-repeat-n: 2
+  --inference-logging-step-interval: 1
+  --log-interval: 1
+  --timing-log-level: 0
+
+  # Sampling
+  --temperature: 1.0
+  --top_k: 1
+  --return-log-probs: true
+  --num-tokens-to-generate: 16
+
+  --output-path: ${INFERENCE_OUTPUT_PATH}
+  --prompts: "The capital of France is"
+
+METRICS:
+  - "generated_tokens"
+  - "logprobs"

tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml

@@ -0,0 +1,65 @@
+type: basic
+format_version: 1
+maintainers: [mcore]
+loggers: [stdout]
+spec:
+  name: '{test_case}_{environment}_{platforms}'
+  model: moe
+  build: mcore-pyt-{environment}
+  nodes: 2
+  gpus: 4
+  n_repeat: 1
+  platforms: dgx_gb200
+  script_setup: |
+    set -euo pipefail
+    unset https_proxy
+    echo "machine gitlab-master.nvidia.com login okoenig password $RO_API_TOKEN" | tee -a /root/.netrc
+
+    # Checkout latest
+    cd /opt
+    rm -rf /opt/megatron-lm; mkdir megatron-lm; cd megatron-lm
+    git init
+    git remote add origin $MCORE_REPO
+    git fetch origin '+refs/merge-requests/*:refs/remotes/merge-requests/*'
+    git fetch origin $MCORE_MR_COMMIT
+    git checkout $MCORE_MR_COMMIT
+    git rev-parse HEAD
+    # Checkout backwards-ref
+    cd /opt
+    rm -rf /opt/megatron-lm-legacy; mkdir megatron-lm-legacy; cd megatron-lm-legacy
+    git init
+    git remote add origin $MCORE_REPO
+    git fetch origin $MCORE_BACKWARDS_COMMIT
+    git checkout $MCORE_BACKWARDS_COMMIT
+    git rev-parse HEAD
+    rm -rf megatron; cp -a /opt/megatron-lm/megatron ./
+  script: |-
+    set -euo pipefail
+    ls
+    cd /opt/megatron-lm
+    export GPUS_PER_NODE={gpus}
+
+    ARGUMENTS=(
+        "CHECKPOINT_LOAD_PATH=/mnt/artifacts"
+        "CHECKPOINT_SAVE_PATH=/tmp/checkpoints"
+        "DATA_PATH=null"
+        "DATA_CACHE_PATH=/workspace/data/cache"
+        "TRAINING_SCRIPT_PATH=examples/inference/advanced/gpt_dynamic_inference.py"
+        "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml"
+        "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json"
+        "OUTPUT_PATH={assets_dir}"
+        "TENSORBOARD_PATH={assets_dir}/tensorboard"
+        "INFERENCE_OUTPUT_PATH={assets_dir}/golden_values_{environment}_{platforms}.json"
+        "N_REPEAT={n_repeat}"
+        "ENABLE_LIGHTWEIGHT_MODE=${{ENABLE_LIGHTWEIGHT_MODE:-}}"
+        "RECORD_CHECKPOINTS=${{RECORD_CHECKPOINTS:-}}"
+    )
+
+    bash ./tests/functional_tests/shell_test_utils/run_ci_test.sh ${{ARGUMENTS[@]}}
+
+products:
+  - test_case: [gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa]
+    products:
+      - environment: [dev]
+        scope: [mr]
+        platforms: [dgx_gb200]

tests/test_utils/recipes/h100/moe-dynamic-inference.yaml

@@ -67,6 +67,11 @@ products:
       - environment: [dev]
         scope: [mr]
         platforms: [dgx_h100]
+  - test_case: [gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa]
+    products:
+      - environment: [dev]
+        scope: [mr]
+        platforms: [dgx_h100]
   - test_case: [gpt_dynamic_inference_tp4_pp1_ep4_16B_prefix_caching]
     products:
       - environment: [dev]