Gemma4e2b jax.checkpoint & remat fix by mattdonati · Pull Request #4847 · AI-Hypercomputer/maxtext

mattdonati · 2026-08-12T01:30:54Z

Description

When training the Gemma 4 Small (E2B/E4B) model with NEW_MODEL_DESIGN=1, XLA LICM was co-scheduling the unrolled 35 decoder layers concurrently, causing the peak HBM overhead during the backward pass to spike past 400GB+ for per device for max_target_length=24576 & per_device_batch_size=8

Before: max_target_length=24576 per_device_batch_size=8

export LIBTPU_INIT_ARGS="\
--xla_tpu_dvfs_p_state=7 \
--xla_tpu_scoped_vmem_limit_kib=65536 \
--xla_tpu_num_sparse_cores_for_gather_offloading=1 \
--xla_tpu_bf16_emission_mode=NATIVE_EMISSION \
--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \
--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \
--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \
--xla_tpu_use_tc_device_shape_on_sc=True \
--xla_sc_disable_megacore_partitioning=True \
--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \
--xla_tpu_prefer_async_allgather_to_allreduce=true \
--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \
--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \
--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \
--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \
--xla_tpu_enable_concurrent_sparse_core_offloading=true \
--xla_tpu_enable_offloading_gather_to_sparsecore=true \
--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \
--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \
--xla_tpu_enable_sparse_core_collective_aggregator=true \
--xla_tpu_enable_latency_hiding_layer_scheduler=true \
--xla_tpu_scheduler_percent_shared_memory_limit=150 \
--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \
--xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \
--xla_tpu_pcie_bandwidth_multiplier=0.03 \
--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \
--xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \
--xla_tpu_enable_3d_reduce_scatter_decomposer=false \
--xla_enable_async_all_gather=true"

export MAXTEXT_COMMAND="export LIBTPU_INIT_ARGS=\"$LIBTPU_INIT_ARGS\" && \
HF_TOKEN='' \
JAX_RANDOM_WEIGHTS=1 \
NEW_MODEL_DESIGN=1 \
TPU_MIN_LOG_LEVEL=0 \
TF_CPP_MIN_LOG_LEVEL=0 \
TPU_STDERR_LOG_LEVEL=0 \
python3 -m maxtext.trainers.pre_train.train \
src/maxtext/configs/base.yml \
model_name=gemma4-e2b \
run_name=$WORKLOAD_NAME \
base_output_directory=$BASE_OUTPUT_DIRECTORY \
steps=10 \
skip_first_n_steps_for_profiler=5 \
profiler_steps=5 \
profiler=xplane \
gradient_clipping_threshold=1.0 \
remat_policy=full \
attention=flash \
max_target_length=24576 \
learning_rate=1e-6 \
per_device_batch_size=8 \
allow_split_physical_axes=True \
checkpoint_period=100 \
max_num_checkpoints_to_keep=1000 \
enable_checkpointing=False \
dataset_type=synthetic \
scan_layers=False \
num_vocab_tiling=8"

xpk workload create \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT_ID}" \
  --tpu-type=tpu7x-8 \
  --zone="${ZONE}" \
  --num-slices=1 \
  --docker-image="${BASE_DOCKER_IMAGE}" \
  --workload="${WORKLOAD_NAME}" \
  --command="${MAXTEXT_COMMAND}"

E0810 15:08:29.958635     808 deepsea_compiler_base.cc:3862] XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 421.33G of 94.74G hbm. Exceeded hbm capacity by 326.59G. Memory usage report: 

Memory Space: default (color=0)
Total bytes: 452401005824 (421.33GiB)
  cumulative_size;       size;       offset; used_by_n_values; shapes_list
  ------------------------------------------------------------
   12.00GiB(  2%);   12.00GiB;  12884901888;                1; bf16[24576,262144]
   24.00GiB(  4%);   12.00GiB;            0;              176; bf16[1572864,1120], bf16[256,8,24576,8], bf16[262144,8960], bf16[8,1,24576,1536], bf16[8,24576,8960], bf16[8,8,24576,1120], bf16[8,8,24576,512], 2×bf16[24576,262144], 2×bf16[8,24576,1536], 2×bf16[8,24576,35,256], 2×bf16[8,8,24576,256], 3×bf16[8,24576,8,512], 5×bf16[8,24576,8,256], 14×bf16[8,24576,6144], 20×bf16[8,24576,12288], 119×token[]
   28.50GiB(  5%);    4.50GiB; 122406567936;                2; 2×bf16[8,24576,12288]
   33.00GiB(  6%);    4.50GiB;  25769803776;                1; bf16[8,24576,12288]
   37.50GiB(  6%);    4.50GiB; 127238406144;                2; 2×bf16[8,24576,12288]
   42.00GiB(  7%);    4.50GiB;  30601641984;                1; bf16[8,24576,12288]
   46.50GiB(  8%);    4.50GiB; 132070244352;                2; 2×bf16[8,24576,12288]
   51.00GiB(  9%);    4.50GiB;  35433480192;                1; bf16[8,24576,12288]
   55.50GiB( 10%);    4.50GiB; 136902082560;                2; 2×bf16[8,24576,12288]
   60.00GiB( 10%);    4.50GiB;  40265318400;                1; bf16[8,24576,12288]
   64.50GiB( 11%);    4.50GiB; 141733920768;                2; 2×bf16[8,24576,12288]
   69.00GiB( 12%);    4.50GiB;  45097156608;                1; bf16[8,24576,12288]
   73.50GiB( 13%);    4.50GiB;  49928994816;                1; bf16[8,24576,12288]
   78.00GiB( 13%);    4.50GiB;  54760833024;                1; bf16[8,24576,12288]
   82.50GiB( 14%);    4.50GiB; 156229435392;                2; 2×bf16[8,24576,12288]
   87.00GiB( 15%);    4.50GiB;  59592671232;                1; bf16[8,24576,12288]
   91.50GiB( 16%);    4.50GiB; 161061273600;                2; 2×bf16[8,24576,12288]
   96.00GiB( 16%);    4.50GiB;  64424509440;                1; bf16[8,24576,12288]
  100.50GiB( 17%);    4.50GiB;  69256347648;                1; bf16[8,24576,12288]
  105.00GiB( 18%);    4.50GiB;  74088185856;                1; bf16[8,24576,12288]
  109.50GiB( 19%);    4.50GiB; 175556788224;                2; 2×bf16[8,24576,12288]
  114.00GiB( 20%);    4.50GiB;  78920024064;                1; bf16[8,24576,12288]
  118.50GiB( 20%);    4.50GiB; 180388626432;                2; 2×bf16[8,24576,12288]
  123.00GiB( 21%);    4.50GiB;  83751862272;                1; bf16[8,24576,12288]
  127.50GiB( 22%);    4.50GiB; 185220464640;                2; 2×bf16[8,24576,12288]
  132.00GiB( 23%);    4.50GiB;  88583700480;                1; bf16[8,24576,12288]
  136.50GiB( 23%);    4.50GiB;  93415538688;                1; bf16[8,24576,12288]
  141.00GiB( 24%);    4.50GiB; 194884141056;                2; 2×bf16[8,24576,12288]
  145.50GiB( 25%);    4.50GiB;  98247376896;                1; bf16[8,24576,12288]
  150.00GiB( 26%);    4.50GiB; 199715979264;                2; 2×bf16[8,24576,12288]
  154.50GiB( 26%);    4.50GiB; 103079215104;                1; bf16[8,24576,12288]
  159.00GiB( 27%);    4.50GiB; 204547817472;                2; 2×bf16[8,24576,12288]
  163.50GiB( 28%);    4.50GiB; 107911053312;                1; bf16[8,24576,12288]
  168.00GiB( 29%);    4.50GiB; 209379655680;                2; 2×bf16[8,24576,12288]
  172.50GiB( 30%);    4.50GiB; 112742891520;                1; bf16[8,24576,12288]
  177.00GiB( 30%);    4.50GiB; 214211493888;                3; f32[262144,1536], 2×bf16[8,24576,12288]
  181.50GiB( 31%);    4.50GiB; 190052302848;                3; bf16[8,8,24576,512], 2×bf16[8,24576,12288]
  186.00GiB( 32%);    4.50GiB; 170724950016;                3; bf16[8,24576,8,256], 2×bf16[8,24576,12288]
  190.50GiB( 33%);    4.50GiB; 165893111808;                3; bf16[8,24576,1536], 2×bf16[8,24576,12288]
  195.00GiB( 33%);    4.50GiB; 151397597184;                3; bf16[8,24576,8,256], 2×bf16[8,24576,12288]
  199.50GiB( 34%);    4.50GiB; 146565758976;                5; bf16[8,8,24576,256], 2×bf16[8,24576,12288], 2×bf16[8,24576,8,256]
  204.00GiB( 35%);    4.50GiB;   9663676416;               40; 40×bf16[8,24576,12288]
  208.50GiB( 36%);    4.50GiB;   4831838208;               40; 4×bf16[8,24576,1536], 16×bf16[8,24576,6144], 20×bf16[8,24576,12288]
  213.00GiB( 37%);    4.50GiB; 117574729728;                4; bf16[8,24576,1,256], bf16[8,24576,12288], 2×bf16[8,24576,8,512]
  217.38GiB( 37%);    4.38GiB;   4697620480;                2; 2×bf16[262144,8960]
  220.75GiB( 38%);    3.38GiB;   3623878656;                1; bf16[1572864,1120]
  224.03GiB( 38%);    3.28GiB;   9395240960;                2; bf16[24576,8,35,256], bf16[8,24576,35,256]
  227.31GiB( 39%);    3.28GiB;   7046430720;                2; bf16[8,24576,8960], bf16[8,8,24576,1120]
  230.59GiB( 40%);    3.28GiB; 219043332096;                3; bf16[8,24576,35,256], bf16[8,24576,8,256], bf16[8,24576,8,512]
  233.88GiB( 40%);    3.28GiB;   3523215360;                4; bf16[8,24576,6144], bf16[8,24576,8,512], bf16[8,24576,8960], bf16[8,8,24576,1120]
  237.16GiB( 41%);    3.28GiB;   4026531840;                1; bf16[8,24576,35,256]
  240.44GiB( 41%);    3.28GiB;  12918456320;                4; bf16[8,24576,35,256], bf16[8,24576,8,512], 2×bf16[8,24576,8,256]
  242.69GiB( 42%);    2.25GiB; 261221253120;                2; 2×bf16[8,24576,6144]
  244.94GiB( 42%);    2.25GiB; 222566547456;                1; bf16[8,24576,6144]
  247.19GiB( 42%);    2.25GiB; 263637172224;                2; 2×bf16[8,24576,6144]
  249.44GiB( 43%);    2.25GiB; 224982466560;                3; bf16[8,24576,6144], bf16[8,24576,8,256], f32[8,8,24576,128]
  251.69GiB( 43%);    2.25GiB; 266053091328;                2; 2×bf16[8,24576,6144]
  253.94GiB( 44%);    2.25GiB; 227398385664;                1; bf16[8,24576,6144]
  256.19GiB( 44%);    2.25GiB; 268469010432;                2; 2×bf16[8,24576,6144]
  258.44GiB( 44%);    2.25GiB; 229814304768;                1; bf16[8,24576,6144]
  260.69GiB( 45%);    2.25GiB; 270884929536;                2; 2×bf16[8,24576,6144]
  262.94GiB( 45%);    2.25GiB; 232230223872;                2; bf16[8,24576,6144], bf16[8,24576,8,512]
  265.19GiB( 45%);    2.25GiB; 273300848640;                2; 2×bf16[8,24576,6144]
  267.44GiB( 46%);    2.25GiB; 234646142976;                1; bf16[8,24576,6144]
  269.69GiB( 46%);    2.25GiB; 275716767744;                2; 2×bf16[8,24576,6144]
  271.94GiB( 47%);    2.25GiB; 239477981184;                1; bf16[8,24576,6144]
  274.19GiB( 47%);    2.25GiB; 278132686848;                2; 2×bf16[8,24576,6144]
  276.44GiB( 47%);    2.25GiB; 241893900288;                1; bf16[8,24576,6144]
  278.69GiB( 48%);    2.25GiB; 282964525056;                2; 2×bf16[8,24576,6144]
  280.94GiB( 48%);    2.25GiB; 285380444160;                2; 2×bf16[8,24576,6144]
  283.19GiB( 49%);    2.25GiB; 287796363264;                2; 2×bf16[8,24576,6144]
  285.44GiB( 49%);    2.25GiB; 292628201472;                6; bf16[8,8,24576,512], 2×bf16[8,24576,6144], 3×bf16[8,24576,8,512]
  287.69GiB( 49%);    2.25GiB; 290212282368;                3; bf16[8,24576,1536], 2×bf16[8,24576,6144]
  289.94GiB( 50%);    2.25GiB; 256389414912;                8; bf16[8,24576,1536], bf16[8,24576,8,256], bf16[8,24576,8,512], bf16[8,8,24576,512], 4×bf16[8,24576,6144]
  292.19GiB( 50%);    2.25GiB; 251557576704;                4; 4×bf16[8,24576,6144]
  294.44GiB( 50%);    2.25GiB; 280548605952;                3; bf16[8,8,24576,512], 2×bf16[8,24576,6144]
  296.69GiB( 51%);    2.25GiB; 249141657600;                5; bf16[8,24576,8,512], 4×bf16[8,24576,6144]
  298.94GiB( 51%);    2.25GiB; 244309819392;                3; bf16[8,24576,8,512], 2×bf16[8,24576,6144]
  301.19GiB( 52%);    2.25GiB; 246725738496;                4; 4×bf16[8,24576,6144]
  303.44GiB( 52%);    2.25GiB; 237062062080;                2; bf16[8,24576,256], bf16[8,24576,6144]
  305.69GiB( 52%);    2.25GiB; 253973495808;                6; bf16[8,24576,8,256], bf16[8,8,24576,512], 4×bf16[8,24576,6144]
  307.94GiB( 53%);    2.25GiB; 258805334016;               10; bf16[8,24576,1536,1], 2×bf16[8,24576,6144], 7×bf16[8,24576,1536]
  310.19GiB( 53%);    2.25GiB;   2415919104;               12; 2×bf16[8,24576,1536], 10×bf16[8,24576,6144]
  312.44GiB( 54%);    2.25GiB;   8355053568;                2; bf16[8,24576,6144], bf16[8,24576,8,256]
  314.69GiB( 54%);    2.25GiB;   5939134464;                2; 2×bf16[8,24576,6144]
  316.19GiB( 54%);    1.50GiB; 231928233984;                1; bf16[8,24576,8,512]
  317.69GiB( 54%);    1.50GiB;   6744440832;                1; bf16[8,24576,8,512]
  319.19GiB( 55%);    1.50GiB;   5133828096;                1; bf16[8,24576,8,512]
  320.69GiB( 55%);    1.50GiB; 223875170304;                1; bf16[8,24576,8,512]
  322.19GiB( 55%);    1.50GiB;  14529069056;                1; bf16[8,24576,8,512]
  323.69GiB( 55%);    1.50GiB; 228707008512;                1; bf16[8,24576,8,512]
  325.19GiB( 56%);    1.50GiB; 220653944832;                2; bf16[8,24576,8,256], bf16[8,24576,8,512]
  326.69GiB( 56%);    1.50GiB; 230317621248;                2; bf16[8,24576,8,256], bf16[8,24576,8,512]
  328.19GiB( 56%);    1.50GiB; 222264557568;                1; bf16[8,24576,8,512]
  329.69GiB( 57%);    1.50GiB; 225485783040;                2; bf16[8,24576,8,256], bf16[8,24576,8,512]
  331.19GiB( 57%);    1.50GiB; 227096395776;                1; bf16[8,24576,8,512]
  332.69GiB( 57%);    1.50GiB; 233840836608;                1; bf16[8,24576,8,512]
  334.19GiB( 57%);    1.50GiB; 295044120576;                1; bf16[8,8,24576,512]
  335.69GiB( 58%);    1.50GiB; 296654733312;                2; bf16[8,24576,8,512], bf16[8,8,24576,512]
  337.19GiB( 58%);    1.50GiB; 298265346048;                2; bf16[8,24576,8,512], bf16[8,8,24576,512]
  338.69GiB( 58%);    1.50GiB; 299875958784;                2; bf16[8,24576,8,512], bf16[8,8,24576,512]
  340.19GiB( 58%);    1.50GiB; 303097184256;               12; bf16[8,8,24576,512], 2×bf16[8,24576,8,512], 4×bf16[8,24576,8,256], 5×f32[8,8,24576,128]
  341.69GiB( 59%);    1.50GiB; 304707796992;               13; bf16[8,8,24576,512], 2×bf16[8,24576,8,512], 5×bf16[8,24576,8,256], 5×f32[8,8,24576,128]
  343.19GiB( 59%);    1.50GiB; 119185342464;                1; bf16[512,8,24576,8]
  344.69GiB( 59%);    1.50GiB; 215822106624;                4; bf16[8,24576,1536], bf16[8,8,24576,512], 2×bf16[8,24576,8,512]
  346.19GiB( 59%);    1.50GiB; 315982086144;               10; bf16[256,8,24576,8], bf16[8,24576,1536,1], bf16[8,24576,8,128], bf16[8,8,24576,512], 2×bf16[8,24576,8,256], 4×bf16[8,24576,1536]
  347.69GiB( 60%);    1.50GiB; 191662915584;                1; bf16[8,24576,8,512]
  349.19GiB( 60%);    1.50GiB; 193273528320;                1; bf16[512,8,24576,8]
  350.69GiB( 60%);    1.50GiB; 217432719360;               30; 2×bf16[8,24576,8,512], 2×bf16[8,8,24576,512], 4×bf16[8,24576,1536,1], 4×bf16[8,8,24576,256], 8×bf16[8,24576,8,256], 10×bf16[8,24576,1536]
  352.19GiB( 60%);    1.50GiB; 314371473408;               12; bf16[8,24576,1536,1], bf16[8,8,24576,512], 2×bf16[8,8,24576,256], 4×bf16[8,24576,1536], 4×bf16[8,24576,8,256]
  353.69GiB( 61%);    1.50GiB;  16441671680;                1; bf16[8,24576,8,512]
  355.19GiB( 61%);    1.50GiB; 312760860672;                7; bf16[8,8,24576,256], bf16[8,8,24576,512], 2×bf16[8,24576,1536], 3×bf16[8,24576,8,256]
  356.69GiB( 61%);    1.50GiB; 301486571520;                6; bf16[8,24576,8,256], bf16[8,8,24576,512], 2×bf16[8,24576,8,512], 2×f32[8,8,24576,128]
  358.19GiB( 61%);    1.50GiB;  21273509888;                2; bf16[8,24576,8,256], bf16[8,24576,8,512]
  359.69GiB( 62%);    1.50GiB;  22884122624;                2; bf16[512,8,24576,8], bf16[8,24576,8,256]
  361.19GiB( 62%);    1.50GiB;  18052284416;                9; bf16[512,8,24576,8], bf16[8,8,24576,512], 2×bf16[8,24576,8,512], 2×bf16[8,8,24576,256], 3×bf16[8,24576,8,256]
  362.69GiB( 62%);    1.50GiB; 311150247936;               11; bf16[8,24576,1536,1], bf16[8,8,24576,512], 2×bf16[8,24576,1536], 2×bf16[8,8,24576,256], 5×bf16[8,24576,8,256]
  364.19GiB( 62%);    1.50GiB;   1610612736;                6; bf16[512,8,24576,8], bf16[8,8,24576,512], 2×bf16[256,8,24576,8], 2×bf16[8,24576,1536]
  365.69GiB( 63%);    1.50GiB; 309539635200;               15; bf16[8,8,24576,512], 2×bf16[8,24576,1536,1], 3×bf16[8,24576,1536], 4×bf16[8,8,24576,256], 5×bf16[8,24576,8,256]
  367.19GiB( 63%);    1.50GiB; 282159218688;                2; bf16[8,24576,8,512], bf16[8,8,24576,512]
  368.69GiB( 63%);    1.50GiB; 245920432128;                1; bf16[512,8,24576,8]
  370.19GiB( 63%);    1.50GiB;  19662897152;                7; bf16[8,24576,256], bf16[8,24576,8,256], bf16[8,8,24576,512], 2×bf16[256,8,24576,8], 2×bf16[8,24576,1536]
  371.69GiB( 64%);    1.50GiB; 307929022464;               15; bf16[256,8,24576,8], bf16[8,8,24576,512], 2×bf16[8,8,24576,256], 3×bf16[8,24576,1536,1], 4×bf16[8,24576,1536], 4×bf16[8,24576,8,256]
  373.19GiB( 64%);    1.50GiB; 252362883072;                3; bf16[8,24576,8,512], 2×bf16[8,24576,8,256]
  374.69GiB( 64%);    1.50GiB; 255584108544;                2; bf16[512,8,24576,8], bf16[8,24576,8,256]
  376.19GiB( 64%);    1.50GiB; 250752270336;                3; bf16[8,8,24576,512], 2×bf16[8,24576,8,512]
  377.69GiB( 65%);    1.50GiB; 306318409728;                6; bf16[8,24576,1536,1], bf16[8,24576,8,256], bf16[8,8,24576,512], 3×bf16[8,24576,1536]
  379.19GiB( 65%);    1.50GiB; 317592698880;               17; bf16[1,8,24576,1536], bf16[256,8,24576,8], bf16[262144,1536], bf16[8,24576,1536,1], f32[8,8,24576,128], 2×bf16[8,24576,8,256], 5×bf16[8,24576,1536], 5×f32[262144,1536]
  380.31GiB( 65%);    1.12GiB; 233538846720;                2; 2×f32[262144,1120]
  381.44GiB( 65%);    1.12GiB; 234746806272;                3; bf16[8,24576,8,256], 2×f32[262144,1120]
  382.56GiB( 66%);    1.12GiB; 235954765824;                3; bf16[262144,1536], 2×f32[262144,1120]
  383.31GiB( 66%);  768.00MiB;   7247757312;                1; bf16[262144,1536]
  384.06GiB( 66%);  768.00MiB; 240383950848;                1; bf16[8,24576,8,256]
  384.81GiB( 66%);  768.00MiB; 239578644480;                1; bf16[8,24576,8,256]
  385.56GiB( 66%);  768.00MiB; 238773338112;                1; bf16[8,24576,8,256]
  386.31GiB( 66%);  768.00MiB; 237968031744;                1; bf16[8,24576,8,256]
  387.06GiB( 66%);  768.00MiB; 237162725376;                1; bf16[8,24576,8,256]
  387.81GiB( 66%);  768.00MiB; 254778802176;                1; bf16[8,24576,8,256]
  388.56GiB( 67%);  768.00MiB; 253168189440;                1; bf16[8,24576,8,256]
  389.3
E0810 15:08:30.247973     808 deepsea_compiler_base.cc:3882] Failed to verify buffer assignment.
I0810 15:08:30.248222    4815 deepsea_compiler_base.cc:4011] BACKEND_PASSES stage duration: 11.52751936825s
I0810 15:08:30.525268       7 deepsea_compiler_base.cc:989] END_TO_END stage duration: 2m49.006841549s
I0810 15:08:30.525986       7 deepsea_compiler_base.cc:990] Note that END_TO_END duration also includes time outside of HLO_PASSES, BACKEND_PASSES, and CODE_GENERATION

Peak HBM: 421.33GiB

Before: max_target_length=8192 per_device_batch_size=1

export LIBTPU_INIT_ARGS="\
--xla_tpu_dvfs_p_state=7 \
--xla_tpu_scoped_vmem_limit_kib=65536 \
--xla_tpu_num_sparse_cores_for_gather_offloading=1 \
--xla_tpu_bf16_emission_mode=NATIVE_EMISSION \
--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \
--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \
--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \
--xla_tpu_use_tc_device_shape_on_sc=True \
--xla_sc_disable_megacore_partitioning=True \
--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \
--xla_tpu_prefer_async_allgather_to_allreduce=true \
--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \
--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \
--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \
--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \
--xla_tpu_enable_concurrent_sparse_core_offloading=true \
--xla_tpu_enable_offloading_gather_to_sparsecore=true \
--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \
--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \
--xla_tpu_enable_sparse_core_collective_aggregator=true \
--xla_tpu_enable_latency_hiding_layer_scheduler=true \
--xla_tpu_scheduler_percent_shared_memory_limit=150 \
--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \
--xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \
--xla_tpu_pcie_bandwidth_multiplier=0.03 \
--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \
--xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \
--xla_tpu_enable_3d_reduce_scatter_decomposer=false \
--xla_enable_async_all_gather=true"

# Safely inject the variables by exporting them inside the container's shell session
export MAXTEXT_COMMAND="export LIBTPU_INIT_ARGS=\"$LIBTPU_INIT_ARGS\" && \
HF_TOKEN='' \
JAX_RANDOM_WEIGHTS=1 \
NEW_MODEL_DESIGN=1 \
TPU_MIN_LOG_LEVEL=0 \
TF_CPP_MIN_LOG_LEVEL=0 \
TPU_STDERR_LOG_LEVEL=0 \
PYTHONPATH=/app/src \
python3 -m maxtext.trainers.pre_train.train \
src/maxtext/configs/base.yml \
model_name=gemma4-e2b \
run_name=$WORKLOAD_NAME \
base_output_directory=$BASE_OUTPUT_DIRECTORY \
steps=10 \
skip_first_n_steps_for_profiler=5 \
profiler_steps=5 \
profiler=xplane \
gradient_clipping_threshold=1.0 \
remat_policy=qkv_proj_offloaded \
attention=flash \
max_target_length=8192 \
learning_rate=1e-6 \
per_device_batch_size=1 \
allow_split_physical_axes=True \
checkpoint_period=100 \
max_num_checkpoints_to_keep=1000 \
enable_checkpointing=False \
dataset_type=synthetic \
scan_layers=False \
num_vocab_tiling=8"

xpk workload create \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT_ID}" \
  --tpu-type=tpu7x-8 \
  --zone="${ZONE}" \
  --num-slices=1 \
  --base-docker-image="${BASE_DOCKER_IMAGE}" \
  --script-dir="/usr/local/google/home/mattdonati/common_files/maxtext" \
  --workload="${WORKLOAD_NAME}" \
  --command="${MAXTEXT_COMMAND}"

Xprof: https://xprof.corp.google.com/memory_profile/mattdonati-7708346007755985034
Peak HBM: 46.61 GiB

Fix:

Modified _apply_gemma4_small_layers in src/maxtext/layers/nnx_decoders.py to securely wrap both the layer's forward pass and its KV extraction (compute_shared_kv) within a jax.checkpoint boundary.
Plumbed this boundary through a native jax.lax.scan(length=1) block armed with skip-simplify-while-loops_trip-count-one, isolating the compiler graph identically to how the standard _scan_global_layer operates in gemma4.py.
This optimization safely forces peak memory back down to ~75-90GB.

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

After: max_target_length=24576 per_device_batch_size=8

export LIBTPU_INIT_ARGS="\
--xla_tpu_dvfs_p_state=7 \
--xla_tpu_scoped_vmem_limit_kib=65536 \
--xla_tpu_num_sparse_cores_for_gather_offloading=1 \
--xla_tpu_bf16_emission_mode=NATIVE_EMISSION \
--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \
--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \
--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \
--xla_tpu_use_tc_device_shape_on_sc=True \
--xla_sc_disable_megacore_partitioning=True \
--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \
--xla_tpu_prefer_async_allgather_to_allreduce=true \
--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \
--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \
--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \
--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \
--xla_tpu_enable_concurrent_sparse_core_offloading=true \
--xla_tpu_enable_offloading_gather_to_sparsecore=true \
--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \
--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \
--xla_tpu_enable_sparse_core_collective_aggregator=true \
--xla_tpu_enable_latency_hiding_layer_scheduler=true \
--xla_tpu_scheduler_percent_shared_memory_limit=150 \
--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \
--xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \
--xla_tpu_pcie_bandwidth_multiplier=0.03 \
--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \
--xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \
--xla_tpu_enable_3d_reduce_scatter_decomposer=false \
--xla_enable_async_all_gather=true"


export MAXTEXT_COMMAND="export LIBTPU_INIT_ARGS=\"$LIBTPU_INIT_ARGS\" && \
HF_TOKEN='' \
JAX_RANDOM_WEIGHTS=1 \
NEW_MODEL_DESIGN=1 \
TPU_MIN_LOG_LEVEL=0 \
TF_CPP_MIN_LOG_LEVEL=0 \
TPU_STDERR_LOG_LEVEL=0 \
PYTHONPATH=/app/src \
python3 -m maxtext.trainers.pre_train.train \
src/maxtext/configs/base.yml \
model_name=gemma4-e2b \
run_name=$WORKLOAD_NAME \
base_output_directory=$BASE_OUTPUT_DIRECTORY \
steps=10 \
skip_first_n_steps_for_profiler=5 \
profiler_steps=5 \
profiler=xplane \
gradient_clipping_threshold=1.0 \
remat_policy=full \
attention=flash \
max_target_length=24576 \
learning_rate=1e-6 \
per_device_batch_size=8 \
allow_split_physical_axes=True \
checkpoint_period=100 \
max_num_checkpoints_to_keep=1000 \
enable_checkpointing=False \
dataset_type=synthetic \
scan_layers=False \
num_vocab_tiling=8"

xpk workload create \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT_ID}" \
  --tpu-type=tpu7x-8 \
  --zone="${ZONE}" \
  --num-slices=1 \
  --base-docker-image="${BASE_DOCKER_IMAGE}" \
  --script-dir="${SCRIPT_DIR}" \
  --workload="${WORKLOAD_NAME}" \
  --command="${MAXTEXT_COMMAND}"

I0811 21:54:44.943716 138304841739136 metric_logger.py:333] To see full metrics 'tensorboard --logdir=gs:
I0811 21:54:55.531340 138304841739136 metric_logger.py:226] completed step: 1, seconds: 0.631, TFLOP/s/device: 5688.282, Tokens/s/device: 311629.521, total_weights: 1572864, loss: 12.979, lm_loss: 12.979, perplexity: 433358.281
I0811 21:55:06.119312 138304841739136 metric_logger.py:226] completed step: 2, seconds: 10.228, TFLOP/s/device: 350.869, Tokens/s/device: 19222.175, total_weights: 1572864, loss: 12.980, lm_loss: 12.980, perplexity: 433463.719
I0811 21:55:16.704861 138304841739136 metric_logger.py:226] completed step: 3, seconds: 10.344, TFLOP/s/device: 346.932, Tokens/s/device: 19006.461, total_weights: 1572864, loss: 12.979, lm_loss: 12.979, perplexity: 433110.281
I0811 21:55:27.300509       7 profiler_session.cc:121] Profiler session initializing.
I0811 21:55:27.300535       7 profiler_session.cc:136] Profiler session started.
I0811 21:55:27.300967       7 xprof_profiler.cc:340] CreateGrpcTslProfiler with host_tracer_level=2
I0811 21:55:27.377263     769 tpu_vxc_driver.cc:3016] ReadGlobalTimeCounter returning real data
W0811 21:55:27.378361    4942 gfc_configs.cc:136]  Overriding tracemark filters for GFC to 0..max. See b/394884361.
W0811 21:55:27.378389    4942 gfc_configs.cc:136]  Overriding tracemark filters for GFC to 0..max. See b/394884361.
W0811 21:55:27.378797     769 gfc_configs.cc:136]  Overriding tracemark filters for GFC to 0..max. See b/394884361.
W0811 21:55:27.378809     769 gfc_configs.cc:136]  Overriding tracemark filters for GFC to 0..max. See b/394884361.
I0811 21:55:27.414536     769 latent_see.cc:155] Latent-see collection enabling
I0811 21:55:27.440104 138304841739136 metric_logger.py:226] completed step: 4, seconds: 10.588, TFLOP/s/device: 338.930, Tokens/s/device: 18568.092, total_weights: 1572864, loss: 12.979, lm_loss: 12.979, perplexity: 433246.562
I0811 21:55:28.014659     769 latent_see.cc:159] Latent-see collection recording
I0811 21:55:38.016342 138304841739136 metric_logger.py:226] completed step: 5, seconds: 21.296, TFLOP/s/device: 168.519, Tokens/s/device: 9232.233, total_weights: 1572864, loss: 12.979, lm_loss: 12.979, perplexity: 433426.469
I0811 21:55:48.608531 138304841739136 metric_logger.py:226] completed step: 6, seconds: 0.880, TFLOP/s/device: 4080.193, Tokens/s/device: 223531.217, total_weights: 1572864, loss: 12.979, lm_loss: 12.979, perplexity: 433264.312
I0811 21:55:59.200188 138304841739136 metric_logger.py:226] completed step: 7, seconds: 9.733, TFLOP/s/device: 368.706, Tokens/s/device: 20199.361, total_weights: 1572864, loss: 12.979, lm_loss: 12.979, perplexity: 433307.281
I0811 21:56:20.393859       7 profiler_session.cc:84] Profiler session collecting data.
I0811 21:56:20.926181     769 latent_see.cc:169] Latent-see collection stopping
I0811 21:56:20.926638     769 latent_see.cc:175] Latent-see collection stopped: processing 61 bins
I0811 21:56:20.927149     769 latent_see.cc:237] Latent-see collection complete
I0811 21:56:24.033617       7 timer.cc:51] ConvertTpuTraceToXPlane: 2467.63 ms (elapsed)
W0811 21:56:24.465452     769 runtime_metric_aggregator.cc:104] Failed to get metric 'tpu.runtime.hbm.memory.total.bytes' from worker localhost:8432: failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:8432: Failed to connect to remote host: Connection refused
I0811 21:56:25.046963       7 xplane_hlo_fixer.cc:117] XProf HLO fixer triggered for metadata plane.
I0811 21:56:29.503304       7 save_profile.cc:239] Collecting XSpace to repository: gs://
I0811 21:56:38.051955       7 save_profile.cc:212] Creating directory: gs://
I0811 21:56:44.472283       7 save_profile.cc:218] Dumped gzipped tool data for trace.json.gz to gs:
tnn8.trace.json.gz
I0811 21:56:46.834911       7 profiler_session.cc:167] Profiler session tear down.
I0811 21:56:46.839565 138304841739136 metric_logger.py:226] completed step: 8, seconds: 10.593, TFLOP/s/device: 338.777, Tokens/s/device: 18559.733, total_weights: 1572864, loss: 12.979, lm_loss: 12.979, perplexity: 433310.156
I0811 21:56:46.850301 138304841739136 metric_logger.py:226] completed step: 9, seconds: 10.590, TFLOP/s/device: 338.893, Tokens/s/device: 18566.051, total_weights: 1572864, loss: 12.980, lm_loss: 12.980, perplexity: 433484.906
Per train step:
 Total TFLOPs: 3588.75 
 split as 74.93% learnable weight flops and 25.07% attention flops

Peak HBM: 82.94 GiB
Xprof: http://xprof.corp.google.com/trace_viewer/mattdonati-17081265437688384004

After: max_target_length=8192 per_device_batch_size=1

export LIBTPU_INIT_ARGS="\
--xla_tpu_dvfs_p_state=7 \
--xla_tpu_scoped_vmem_limit_kib=65536 \
--xla_tpu_num_sparse_cores_for_gather_offloading=1 \
--xla_tpu_bf16_emission_mode=NATIVE_EMISSION \
--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \
--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \
--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \
--xla_tpu_use_tc_device_shape_on_sc=True \
--xla_sc_disable_megacore_partitioning=True \
--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \
--xla_tpu_prefer_async_allgather_to_allreduce=true \
--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \
--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \
--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \
--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \
--xla_tpu_enable_concurrent_sparse_core_offloading=true \
--xla_tpu_enable_offloading_gather_to_sparsecore=true \
--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \
--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \
--xla_tpu_enable_sparse_core_collective_aggregator=true \
--xla_tpu_enable_latency_hiding_layer_scheduler=true \
--xla_tpu_scheduler_percent_shared_memory_limit=150 \
--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \
--xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \
--xla_tpu_pcie_bandwidth_multiplier=0.03 \
--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \
--xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \
--xla_tpu_enable_3d_reduce_scatter_decomposer=false \
--xla_enable_async_all_gather=true"

export MAXTEXT_COMMAND="export LIBTPU_INIT_ARGS=\"$LIBTPU_INIT_ARGS\" && \
HF_TOKEN='' \
JAX_RANDOM_WEIGHTS=1 \
NEW_MODEL_DESIGN=1 \
TPU_MIN_LOG_LEVEL=0 \
TF_CPP_MIN_LOG_LEVEL=0 \
TPU_STDERR_LOG_LEVEL=0 \
PYTHONPATH=/app/src \
python3 -m maxtext.trainers.pre_train.train \
src/maxtext/configs/base.yml \
model_name=gemma4-e2b \
run_name=$WORKLOAD_NAME \
base_output_directory=$BASE_OUTPUT_DIRECTORY \
steps=10 \
skip_first_n_steps_for_profiler=5 \
profiler_steps=5 \
profiler=xplane \
gradient_clipping_threshold=1.0 \
remat_policy=qkv_proj_offloaded \
attention=flash \
max_target_length=8192 \
learning_rate=1e-6 \
per_device_batch_size=1 \
allow_split_physical_axes=True \
checkpoint_period=100 \
max_num_checkpoints_to_keep=1000 \
enable_checkpointing=False \
dataset_type=synthetic \
scan_layers=False \
num_vocab_tiling=8"

xpk workload create \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT_ID}" \
  --tpu-type=tpu7x-8 \
  --zone="${ZONE}" \
  --num-slices=1 \
  --base-docker-image="${BASE_DOCKER_IMAGE}" \
  --script-dir="/usr/local/google/home/mattdonati/common_files/maxtext" \
  --workload="${WORKLOAD_NAME}" \
  --command="${MAXTEXT_COMMAND}"

Peak HBM: 17.73 GiB
Xprof: https://xprof.corp.google.com/memory_profile/mattdonati-3052316813571309926

logits check

export WORKLOAD_NAME=logit-check
export CHECKPOINT_DIR="$BUCKET"
export MAXTEXT_COMMAND="pip install torch --index-url https://download.pytorch.org/whl/cpu && \
export PYTHONPATH=/app/src && \
python3 -m tests.utils.forward_pass_logit_checker src/maxtext/configs/base.yml \
model_name=gemma4-e2b \
tokenizer_path=google/gemma-4-E2B \
hf_access_token=\${HF_TOKEN} \
load_parameters_path=${CHECKPOINT_DIR} \
--run_hf_model=true \
--hf_model_path=google/gemma-4-E2B \
--max_kl_div=0.03 \
use_multimodal=false \
scan_layers=false \
per_device_batch_size=1 \
dtype=float32 \
attention=dot_product"

xpk workload create \
--cluster="${CLUSTER_NAME}" \
--project="${PROJECT_ID}" \
--tpu-type=tpu7x-8 \
--zone="${ZONE}" \
--num-slices=1 \
--base-docker-image="${BASE_DOCKER_IMAGE}" \
--script-dir="${SCRIPT_DIR" \
--workload="${WORKLOAD_NAME}" \
--env="HF_TOKEN=${HF_TOKEN}" \
--command="${MAXTEXT_COMMAND}"

--- Prompt: I love to ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 4301       | travel               | 20.9291    |
| 1676       | read                 | 20.9070    |
| 3944       | cook                 | 20.6754    |
| 1386       | make                 | 20.4672    |
| 1460       | see                  | 20.2576    |
| 4903       | write                | 20.0967    |
| 48301      | bake                 | 20.0017    |
| 2619       | create               | 19.9499    |
| 577        | be                   | 19.9396    |
| 1441       | play                 | 19.9361    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 4301       | travel               | 20.9306    |
| 1676       | read                 | 20.9103    |
| 3944       | cook                 | 20.6729    |
| 1386       | make                 | 20.4571    |
| 1460       | see                  | 20.2578    |
| 4903       | write                | 20.1036    |
| 48301      | bake                 | 19.9989    |
| 577        | be                   | 19.9489    |
| 2619       | create               | 19.9458    |
| 1441       | play                 | 19.9446    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 80.0                 |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 1.69e-05
INFO:absl:Per-token KL Divergences: 
['8.10e-06', '1.17e-05', '2.90e-05', '1.87e-05']
INFO:absl:
Max KL divergence for a single token in the set: 2.90e-05
INFO:absl:
--- Prompt: Today is a ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1719       | day                  | 19.7820    |
| 2803       | special              | 19.6162    |
| 1401       | very                 | 19.4278    |
| 2563       | big                  | 19.2898    |
| 1535       | good                 | 18.8568    |
| 1822       | great                | 18.7710    |
| 11019      | sad                  | 17.9225    |
| 861        | new                  | 17.4601    |
| 4148       | beautiful            | 17.2809    |
| 10940      | holiday              | 17.1711    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1719       | day                  | 19.7681    |
| 2803       | special              | 19.6118    |
| 1401       | very                 | 19.4198    |
| 2563       | big                  | 19.2726    |
| 1535       | good                 | 18.8509    |
| 1822       | great                | 18.7621    |
| 11019      | sad                  | 17.9143    |
| 861        | new                  | 17.4404    |
| 4148       | beautiful            | 17.2801    |
| 10940      | holiday              | 17.1665    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 100.0                |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 1.49e-05
INFO:absl:Per-token KL Divergences: 
['8.10e-06', '1.91e-05', '1.45e-05', '1.79e-05']
INFO:absl:
Max KL divergence for a single token in the set: 1.91e-05
INFO:absl:
--- Prompt: What is the ---
INFO:absl:
--- MaxText model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1791       | best                 | 13.7273    |
| 4954       | difference           | 13.7177    |
| 1346       | most                 | 12.8756    |
| 5708       | purpose              | 12.7172    |
| 1689       | main                 | 12.2311    |
| 1463       | name                 | 12.0886    |
| 6590       | meaning              | 12.0761    |
| 5905       | primary              | 11.9940    |
| 3853       | role                 | 11.9537    |
| 1171       | first                | 11.8979    |

INFO:absl:
--- HF model top 10 tokens ---
INFO:absl:| Token ID   | Token                | Score      |
|------------|----------------------|------------|
| 1791       | best                 | 13.7341    |
| 4954       | difference           | 13.7253    |
| 1346       | most                 | 12.8830    |
| 5708       | purpose              | 12.7114    |
| 1689       | main                 | 12.2366    |
| 1463       | name                 | 12.0971    |
| 6590       | meaning              | 12.0800    |
| 5905       | primary              | 11.9974    |
| 3853       | role                 | 11.9618    |
| 1171       | first                | 11.9010    |

INFO:absl:
--- Similarity Metrics of Top Tokens ---
INFO:absl:| Metric                         | Value                |
|--------------------------------|----------------------|
| overlap_count                  | 10/10                |
| jaccard_similarity             | 1.0                  |
| rank_agreement_percentage      | 100.0                |

INFO:absl:
Average KL divergence per token (D_KL(P_golden || Q_model)): 1.86e-05
INFO:absl:Per-token KL Divergences: 
['8.10e-06', '1.15e-05', '3.58e-05', '1.89e-05']
INFO:absl:
Max KL divergence for a single token in the set: 3.58e-05
XPK End: Tue Aug 11 23:35:56 UTC 2026
EXIT_CODE=0

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

gemini-code-assist

Code Review

This pull request refactors the layer application in _apply_gemma4_small_layers to use jax.lax.scan and jax.checkpoint for rematerialization. Feedback on these changes includes removing an unused variable extract_donor_kv, moving the inline import of xla_metadata outside of the loop, and hoisting loop-invariant computations (global_remat_policy and prevent_cse) outside of the decoder layers loop to avoid redundant overhead.

mattdonati requested review from NuojCheng, RissyRan, aireenmei, bvandermoon, gagika, gobbleturk, huytransformer, igorts-git, jiangjy1982, parambole, richjames0, shralex, shuningjin and xibinliu as code owners August 12, 2026 01:30

gemini-code-assist Bot reviewed Aug 12, 2026

View reviewed changes

Comment thread src/maxtext/layers/nnx_decoders.py Outdated

Comment thread src/maxtext/layers/nnx_decoders.py Outdated

Comment thread src/maxtext/layers/nnx_decoders.py Outdated

mattdonati force-pushed the gemma4e2b_checkpoint_fix branch from 6799293 to fc24032 Compare August 12, 2026 15:29

add jax.checkpoint boundaries ot gemma4e2b

46e7797

mattdonati force-pushed the gemma4e2b_checkpoint_fix branch from fc24032 to 46e7797 Compare August 12, 2026 18:13

Navigation Menu

Uh oh!

Gemma4e2b jax.checkpoint & remat fix - #4847

Gemma4e2b jax.checkpoint & remat fix#4847
mattdonati wants to merge 1 commit into
AI-Hypercomputer:mainfrom
mattdonati:gemma4e2b_checkpoint_fix

mattdonati commented Aug 12, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

mattdonati commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

mattdonati commented Aug 12, 2026 •

edited

Loading