You can deploy this model and start earning money today!

MoE-CUDA-Graph-Analysis/
blackwell_run_scripts_echo_disabled

Views

No views yet

Private Megatron-MoE-Echo driver flow

This driver branch is wired to use the private repository:

https://github.com/sriram-17-17/my-Megatron-MoE-Echo.git

with two supported Megatron branches:

moe_hot_expert_poc → referred to as vanilla
moe_hot_expert_poc_hacky_stuffs → referred to as hack

The hack branch disables echo expert cloning. This driver makes that visible in run naming:

vanilla runs/logs/checkpoints use suffix: vanilla_echo_clone
hack runs/logs/checkpoints use suffix: hack_no_echo_clone

1. Setup driver branch

From the driver repository:

bash

cd blackwell_run_scripts
git checkout private-megatron-echo-driver
chmod +x blackwell_run_scripts/*.sh

2. Fresh environment setup

To clone the private Megatron repo and install all dependencies:

Vanilla branch

bash new_setup.sh vanilla

Hack branch

bash new_setup.sh hack

Credential note

The Megatron repo is private. During clone/fetch, Git should prompt for credentials.

Use:

Username: your GitHub username
Password: your GitHub Personal Access Token (PAT)

3. Prepare dataset

bash new_prepare.sh

This creates the GPT2-tokenized Wikitext-103 Megatron dataset under:

/workspace/data/wikitext

4. Run vanilla experiment

If you want a clean vanilla run:

bash

bash change_megatron_echo.sh vanilla
bash clean_checkpoints.sh vanilla
CLEAN_CHECKPOINTS=1 bash train_mixtral_8x_7b.sh

Outputs:

checkpoints: /workspace/checkpoints_mixtral_8x7b_8gpu_vanilla_echo_clone
logs: /workspace/logs/mixtral_8x7b_8gpu_vanilla_echo_clone_<timestamp>.log

5. Run hack experiment

If you want the hacky branch that disables echo expert cloning:

bash

bash change_megatron_echo.sh hack
bash clean_checkpoints.sh hack
CLEAN_CHECKPOINTS=1 bash train_mixtral_8x_7b.sh

Outputs:

checkpoints: /workspace/checkpoints_mixtral_8x7b_8gpu_hack_no_echo_clone
logs: /workspace/logs/mixtral_8x7b_8gpu_hack_no_echo_clone_<timestamp>.log

6. Switching between vanilla and hack

You can switch the local private Megatron checkout at any time:

bash change_megatron_echo.sh vanilla

bash change_megatron_echo.sh hack

This script:

fetches latest changes from the private origin
checks out the requested branch
pulls the branch
updates submodules
reinstalls the checkout with pip install -e .

7. Checkpoint cleanup options

Dedicated cleanup script

bash

bash clean_checkpoints.sh vanilla
bash clean_checkpoints.sh hack
bash clean_checkpoints.sh auto

auto resolves from the currently checked-out Megatron branch.

Cleanup inline at train launch

CLEAN_CHECKPOINTS=1 bash train_mixtral_8x_7b.sh

Also clear old logs for the same mode

CLEAN_LOGS=1 bash clean_checkpoints.sh hack

CLEAN_LOGS=1 CLEAN_CHECKPOINTS=1 bash train_mixtral_8x_7b.sh

8. Resume behavior

The train script only adds --load <checkpoint_dir> if a prior checkpoint actually exists.

That means:

fresh run after cleanup → starts from scratch
rerun without cleanup → resumes from existing checkpoint directory

9. CUDA graph capture workaround

If you hit a CUDA graph capture error during training, replace:

bash

CUDA_GRAPH_ARGS=(
  --enable-cuda-graph
  --cuda-graph-scope full_iteration
  --te-rng-tracker
  --no-check-for-nan-in-loss-and-grad
)

with:

bash

CUDA_GRAPH_ARGS=(
  --external-cuda-graph
  --cuda-graph-scope attn
  --te-rng-tracker
)

in train_mixtral_8x_7b.sh.

10. Minimal end-to-end commands

Vanilla

bash

bash new_setup.sh vanilla
bash new_prepare.sh
bash change_megatron_echo.sh vanilla
bash clean_checkpoints.sh vanilla
CLEAN_CHECKPOINTS=1 bash train_mixtral_8x_7b.sh

Hack

bash

bash new_setup.sh vanilla
bash new_prepare.sh
bash change_megatron_echo.sh hack
bash clean_checkpoints.sh hack
CLEAN_CHECKPOINTS=1 bash train_mixtral_8x_7b.sh
CLEAN_CHECKPOINTS=1 bash train_mixtral_8x_7b_attn_only_cg.sh

MoE-CUDA-Graph-Analysis/blackwell_run_scripts_echo_disabled

Private Megatron-MoE-Echo driver flow

1. Setup driver branch

2. Fresh environment setup

Vanilla branch

Hack branch

Credential note

3. Prepare dataset

4. Run vanilla experiment

5. Run hack experiment

6. Switching between vanilla and hack

7. Checkpoint cleanup options

Dedicated cleanup script

Cleanup inline at train launch

Also clear old logs for the same mode

8. Resume behavior

9. CUDA graph capture workaround

10. Minimal end-to-end commands

Vanilla

Hack

MoE-CUDA-Graph-Analysis/
blackwell_run_scripts_echo_disabled