Update patches + scaling fix v2 + checkpoint restore #61

Draft

yeandy wants to merge 9 commits into

mainfrom

yeandy/update-patches-scaling-patch-v2-checkpoint-restore

Commits on Dec 1, 2025

Update branch
yeandy
committed
Add logging
yeandy
committed

Commits on Dec 2, 2025

Update
yeandy
committed

Commits on Dec 3, 2025

Commits on Jan 30, 2026

Add missing query_pre_attn_scalar
yeandy
committed

Commits on Feb 24, 2026

fix "No ArrayMetadata found" error when enable_single_replica_ckpt_restoring=true
yihuang-amd
committed
fix "stop sending heartbeats" error by jax_distributed_heartbeat_timeout_seconds
yihuang-amd
committed
synchronize hosts before training loop
yihuang-amd
committed