diff --git a/gpt_builders.py b/gpt_builders.py index 57b5179b1a0..a15d672a11d 100644 --- a/gpt_builders.py +++ b/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: diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa/golden_values_dev_dgx_gb200.json new file mode 100644 index 00000000000..4f664c76cf9 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa/golden_values_dev_dgx_h100.json new file mode 100644 index 00000000000..0c6048e989f --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa/model_config.yaml new file mode 100644 index 00000000000..162f9ad6e15 --- /dev/null +++ b/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" diff --git a/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml b/tests/test_utils/recipes/gb200/moe-dynamic-inference.yaml new file mode 100644 index 00000000000..e8728e0b3cb --- /dev/null +++ b/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] diff --git a/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml b/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml index 81255e45d72..828bc15a75a 100644 --- a/tests/test_utils/recipes/h100/moe-dynamic-inference.yaml +++ b/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]