Code

Support CUDA Graph for MoE models#1233

Code

Merged

yaox12 merged 16 commits intoNVIDIA:mainNVIDIA/TransformerEngine:mainfrom buptzyb:cudagraph_moebuptzyb/TransformerEngine:cudagraph_moeCopy head branch name to clipboardNov 25, 2024

+97-17

Commits

Commits on Nov 20, 2024

• Align RNG tracker with megatron

  Show description for 9951d78
  

  buptzyb

  and

  yifeis-nv

  committed
  9951d78

• Fix module_params order and warmup bug in cudagraph

  Show description for b5f7cdf
  

  buptzyb

  and

  yifeis-nv

  committed
  b5f7cdf

• Add fp8_group argument and fix fp8 accuracy issue for cudagraph

  Show description for 0ede5cb
  

  buptzyb

  and

  yifeis-nv

  committed
  0ede5cb

• Add TE modules and weights filters to support MoE models

  Show description for 5596437
  

  buptzyb

  and

  yifeis-nv

  committed
  5596437

• Revert self.fp8

  Show description for 1d0759e
  

  buptzyb

  authored and

  yifeis-nv

  committed
  1d0759e

• Use hooks to filter module params

  Show description for 73a22e2
  

  buptzyb

  authored and

  yifeis-nv

  committed
  73a22e2

• Filter all TE modules in hooks

  Show description for cd56618
  

  buptzyb

  and

  yifeis-nv

  committed
  cd56618

• Format code

  Show description for 2a7f54b
  

  buptzyb

  authored and

  yifeis-nv

  committed
  2a7f54b

• [pre-commit.ci] auto fixes from pre-commit.com hooks

  Show description for c6dddaf
  

  pre-commit-ci[bot]

  authored and

  yifeis-nv

  committed
  c6dddaf

• Update graph.py

  Show description for a01602e
  

  yaox12

  authored and

  yifeis-nv

  committed
  a01602e

• Revert CudaRNGStatesTracker

  Show description for 8017b6d
  

  buptzyb

  authored and

  yifeis-nv

  committed
  8017b6d

• Format Update

  Show description for 41d6100
  yifeis-nv

  committed
  41d6100

• Revert "Use hooks to filter module params"

  Show description for c82faa2
  yifeis-nv

  committed
  c82faa2

• Remove filtering module params

  Show description for 938f325
  buptzyb

  committed
  938f325

Commits on Nov 21, 2024

• Merge branch 'main' into cudagraph_moe

  yaox12

  authored
  1522ecc

Commits on Nov 23, 2024

• Merge branch 'main' into cudagraph_moe

  timmoon10

  authored
  4487315