type: basic
format_version: 1
maintainers: [mcore]
loggers: [stdout]
spec:
  name: "{test_case}_{environment}"
  model: moe
  build: mcore-pyt-{environment}
  nodes: 1
  gpus: 8
  n_repeat: 5
  platforms: dgx_a100
  artifacts:
    /workspace/data/gpt3_data: text/the_pile/shard00
    /workspace/checkpoints/gpt3_mr_mcore_te_tp2_pp1_frozen_resume_torch_dist_te_8experts2parallel_dist_optimizer_dgx_a100_1N8G_dev: model/mcore_gpt/gpt3_mr_mcore_te_tp2_pp1_frozen_resume_torch_dist_te_8experts2parallel_dist_optimizer_dgx_a100_1N8G_dev/22410107
  script: |-
    ls
    cd /opt/megatron-lm

    ARGUMENTS=(
        "DATA_PATH=/workspace/data/gpt3_data"
        "DATA_CACHE_PATH=/workspace/data/cache"
        "OUTPUT_PATH={assets_dir}"
        "TENSORBOARD_PATH={assets_dir}/tensorboard"
        "CHECKPOINT_SAVE_PATH={artifacts_dir}/checkpoints"
        "CHECKPOINT_LOAD_PATH=/workspace/checkpoints/{name}"
        "TRAINING_SCRIPT_PATH=pretrain_gpt.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}.json"
        "N_REPEAT={n_repeat}"
    )

    bash ./tests/functional_tests/shell_test_utils/run_ci_test.sh ${{ARGUMENTS[@]}}

products:
  #######################################################################
  # Nightly tests: Run both DEV and LTS unless something is flaky       #
  #######################################################################
  - test_case: [gpt3_345m_nightly_dgx_a100_1N8G_mcore_tp2_pp2_ep2_te_4experts2parallel]
    products:
      - environment: [dev, lts]
        scope: [nightly]
  - test_case: [gpt3_345m_nightly_dgx_a100_1N8G_mcore_tp2_cp2_pp2_ep2_te_4experts2parallel]
    products:
      - environment: [dev]
        scope: [nightly]
  - test_case: [gpt3_345m_nightly_dgx_a100_1N8G_mcore_tp2_pp2_ep2_etp2_te_4experts2parallel]
    products:
      - environment: [dev, lts]
        scope: [nightly]
  - test_case: [gpt3_345m_nightly_dgx_a100_1N8G_mcore_tp2_pp2_ep2_etp2_te_4experts2parallel_dp_last]
    products:
      - environment: [dev, lts]
        scope: [nightly]
  - test_case: [gpt3_345m_nightly_dgx_a100_1N8G_mcore_tp2_cp2_pp2_ep2_te_4experts2parallel_dp_last]
    products:
      - environment: [dev]
        scope: [nightly]
  - test_case: [gpt3_345m_nightly_dgx_a100_1N8G_mcore_cp2_pp2_ep2_te_4experts2parallel_nondeterministic]
    products:
      - environment: [dev, lts]
        scope: [nightly]
  - test_case: [gpt3_345m_nightly_dgx_a100_1N8G_mcore_cp2_pp2_ep2_te_4experts2parallel_nondeterministic_dp_last]
    products:
      - environment: [dev, lts]
        scope: [nightly]
  - test_case: [gpt3_345m_nightly_dgx_a100_1N8G_tp2_pp2_resume_torch_4experts]
    products:
      - environment: [dev, lts]
        scope: [nightly]
  # - test_case: [gpt3_345m_nightly_dgx_a100_1N8G_mcore_tp2_pp2_resume_torch_dist_te_2experts]
  #   products:
  #     - environment: [dev, lts]
  #       scope: [nightly]
  #######################################################################
  # Weekly tests: Run both DEV and LTS unless something is flaky        #
  #######################################################################
  #######################################################################
  # MR tests: Mostly DEV on MR, and LTS on nightly cadence, except for  #
  #             some very important tests.                              #
  #######################################################################
  - test_case: [gpt3_mr_mcore_te_tp1_pp2_resume_torch_dist_reshard_2x1x4_te_8experts2parallel_dist_optimizer_dgx_a100_1N8G]
    products:
      - environment: [dev]
        scope: [mr]
  - test_case: [gpt3_mr_mcore_te_tp2_pp1_resume_torch_dist_te_8experts2parallel_dist_optimizer_dgx_a100_1N8G]
    products:
      - environment: [dev]
        scope: [mr]
      - environment: [lts]
        scope: [nightly]
  - test_case: [gpt3_mr_mcore_te_tp2_pp1_te_8experts2parallel_overlap_grad_reduce_param_gather_groupedGEMM_dgx_a100_1N8G]
    products:
      - environment: [dev]
        scope: [mr]
      - environment: [lts]
        scope: [nightly]
  - test_case: [gpt3_mr_mcore_te_tp2_pp1_te_8experts_etp1_ep4_dgx_a100_1N8G]
    products:
      - environment: [dev]
        scope: [mr]
      - environment: [lts]
        scope: [nightly]
  # TODO: Enable this test once the compatibility issue is fixed
  # - test_case: [gpt3_mr_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router_dgx_a100_1N8G]
  #   products:
  #     - environment: [dev]
  #       scope: [mr]
  #     - environment: [lts]
  #       scope: [nightly]

  # - test_case: [gpt3_mr_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective_dgx_a100_1N8G]
  #   products:
  #     - environment: [dev]
  #       scope: [mr]
  #     - environment: [lts]
  #       scope: [nightly]
  # - test_case: [gpt3_mr_mcore_te_tp2_pp1_fsdp2_resume_torch_dist_dgx_a100_1N8G]
  #   products:
  #     - environment: [dev]
  #       scope: [mr]
  #     - environment: [lts]
  #       scope: [nightly]
  #######################################################################
  # Super important MR tests that run for both DEV and LTS per MR       #
  #######################################################################

  # Temporarily skipping these tests as both !2719 and !2668 require checkpoint updates. Will update all checkpoints together later.
  # - test_case: [gpt3_mr_mcore_te_tp2_pp1_frozen_resume_torch_dist_te_8experts2parallel_dist_optimizer_dgx_a100_1N8G]
  #   products:
  #     - environment: [dev]
  #       scope: [mr]
  # - test_case: [gpt3_mr_mcore_te_tp2_pp1_frozen_resume_torch_dist_te_8experts2parallel_groupedGEMM_dgx_a100_1N8G]
  #   products:
  #     - environment: [dev]
  #       scope: [mr]
  ###########################
  # Merge train tests       #
  ###########################
  - test_case: [gpt3_moe_mr_mcore_te_ep8_resume_torch_dist_dist_optimizer]
    products:
      - environment: [dev]
        scope: [mr, mr-slim]
  - test_case: [gpt3_moe_mr_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer]
    products:
      - environment: [dev]
        scope: [mr, mr-slim]
