Code

Quantize megablox#1062

Code

Open

lenscloth wants to merge 43 commits intomainAI-Hypercomputer/maxtext:mainfrom quantize_megabloxAI-Hypercomputer/maxtext:quantize_megabloxCopy head branch name to clipboard

+1,088-30

Commits

Commits on Nov 13, 2024

• base

  ZhiyuLi-goog

  committed
  61d944d

• fix import error in pallas megablox kernel

  ZhiyuLi-goog

  committed
  3d25eb3

• Quantize GMM kernel

  ZhiyuLi-goog

  committed
  4a3d9ba

• [MoE] fix typo

  ZhiyuLi-goog

  committed
  2ebe80e

• walk around quantize_params

  ZhiyuLi-goog

  committed
  0e59567

• inferencing shape hack

  ZhiyuLi-goog

  committed
  bd208ed

Commits on Nov 23, 2024

• Enable (1) checkpoint conversion for MoEBlock when megablox=True (2) load quantized weights from AqtEinsum to feed to gmm kernel later.

  lenscloth

  committed
  57e2383

• Fix lint error

  lenscloth

  committed
  9117c40

Commits on Nov 25, 2024

• [rollback] quantize param when ckpt is not quantized.

  lenscloth

  committed
  35e243a

• Merge branch 'quantize_megablox'

  lenscloth

  committed
  94bdc94

Commits on Nov 27, 2024

• Let gmm accept pre-quantized rhs (weight).

  Show description for fa5d5f3
  lenscloth

  committed
  fa5d5f3

Commits on Dec 4, 2024

• Support different quantization precision for lhs / rhs of megablox.

  Show description for 11f1f43
  lenscloth

  committed
  11f1f43

Commits on Dec 5, 2024

• Fix license

  lenscloth

  committed
  1844936

• Refactoring MoEBlock

  lenscloth

  committed
  0d9db17

• Rename in_out_block_spec for better readability

  lenscloth

  committed
  cd8b2d1

Commits on Dec 6, 2024

• Fix lint error & refactor retrieving quantized weight logic

  lenscloth

  committed
  cfd3b15

• fix lint error

  lenscloth

  committed
  5d51030

Commits on Dec 9, 2024

• Merge branch 'quantize_megablox'

  lenscloth

  committed
  3971306

Commits on Dec 14, 2024

• Read lhs and rhs quantization dtype directly from DotGeneral

  lenscloth

  committed
  325c4c3

• Fix MoE related tests

  

  RissyRan

  authored and

  lenscloth

  committed
  95bb7c1

• Add checkpoint topology discovery for the Replicator Service

  

  xuefgu

  authored and

  lenscloth

  committed
  7addc0f

• Fix local restore by re-mapping device ids directly instead of inferring them from how process indexes changed across restarts with some false assumptions.

  Show description for f6c7c5a
  

  cpgaffney1

  authored and

  lenscloth

  committed
  f6c7c5a

• Compact the number of variables for the prefill result cache to reduce Python layer overhead

  

  Zhihao Shan

  authored and

  lenscloth

  committed
  dca3ee5

• add more MoE tests

  

  RissyRan

  authored and

  lenscloth

  committed
  c993cf2

• update setup_gcsfuse for better perf

  

  aireenmei

  authored and

  lenscloth

  committed
  86d0b54

• Assert multiple slices available when requesting DCN parallelisms

  

  gobbleturk

  authored and

  lenscloth

  committed
  346218a

• Add llama 3.1 70b config

  

  raymondzouu

  authored and

  lenscloth

  committed
  f6c38f1

• Update replicator.yaml to include framework and num_slices information

  

  xuefgu

  authored and

  lenscloth

  committed
  7a4a44e

• Support JAX_VERSION for nightly mode on GPU

  

  bvandermoon

  authored and

  lenscloth

  committed
  e3823c0

• Added a custom_wheel mode for building the dependency image.

  Show description for e17fcc6
  

  lukebaumann

  authored and

  lenscloth

  committed
  e17fcc6

• clean up pipeline config setting in its own method

  

  gobbleturk

  authored and

  lenscloth

  committed
  4cdc15d

• Fixes for dropping

  

  mailvijayasingh

  authored and

  lenscloth

  committed
  7390f56

• Fix setup_training_state

  

  khatwanimohit

  authored and

  lenscloth

  committed
  5831547

• point to new jax github location in documentation

  Show description for 90430e1
  

  jakeharmon8

  authored and

  lenscloth

  committed
  90430e1

• Change awk regex command to capture the coordinate address properly

  

  michelle-yooh

  authored and

  lenscloth

  committed
  2d1c51a

• Fixes non-hashable error in ragged attn.

  

  patemotter

  authored and

  lenscloth

  committed
  53a6abe

• Add new remat policy for save_dot_except_mlp with context

  

  bvandermoon

  authored and

  lenscloth

  committed
  4546d68

• Fix moe logging to differentiate dense and megablox runs

  lenscloth

  committed
  945ee3d

• update sharding

  

  ZhiyuLi-goog

  authored and

  lenscloth

  committed
  2b62fe8

• Reshape based on the original input shape

  lenscloth

  committed
  742019f

• Resolve math import error

  lenscloth

  committed
  91b84f6

• Merge branch 'main' into quantize_megablox

  lenscloth

  authored
  cd8e6bf

• fix lint error

  lenscloth

  committed
  b0f21f9