[ROCm] reduce temp compile memory usage before training starts#85

Closed

cj401-amd wants to merge 29 commits into

mainfrom

cj-fix-tmp-mem_rocm-main

Commits on Apr 23, 2026

Commits on Apr 30, 2026

update with a reduction of temp mem usage
Nightly Sweep
committed
upate to reduce temp mem usage

cj401-amd
authored and
Nightly Sweep
committed
update for reduce temp memory usage for ep=2 pp=4

cj401-amd
authored and
Nightly Sweep
committed
Apply deepseek.py temp mem fix from cj-reduce-tmp-mem_rocm-main

Nightly Sweep
and
claude
committed
Fix jax.jit decorator syntax for JAX 0.7.1 compatibility

Nightly Sweep
and
claude
committed
Guard jax_remove_size_one_mesh_axis_from_type for JAX 0.7.1 compat

Nightly Sweep
and
claude
committed
Fix pp=8 ep=1 temp memory regression (36.8 GB -> ~31 GB)

Nightly Sweep
and
claude
committed
Add skip_trivial_specs to pipeline._maybe_shard_with_logical

Nightly Sweep
and
claude
committed
Remove unconditional nan_to_num on all float gradients

Nightly Sweep
and
claude
committed
Clean up debug artifacts before PR

Nightly Sweep
and
claude
committed
Fix depth_scaling regression and attention_op.py TE compatibility issues

Nightly Sweep
and
claude
committed
Fix issues 5/6/7: moe tile comment, embeddings axis, float32_weight_sum sync

Nightly Sweep
and
claude
committed