Skip to content

Backfill CHANGELOG.md + scrub internal-only content in oss_sync.sh (#272) #272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wz337 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Backfill CHANGELOG.md + scrub internal-only content in oss_sync.sh (#272) #272

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
307 changes: 302 additions & 5 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,307 @@
# Changelog

Release notes for the OSS mirror, generated by `scripts/oss_sync.sh` from the internal dev tree.
Release notes for the OSS mirror, generated by `scripts/oss_sync.sh` from the internal source tree.

## 2026-07-24

- [shampoo] Drop experimental `_` prefix + note from lossless distributors — The FullyShard/HybridShard lossless distributors were named
- [shampoo][v3] Rank-balanced per-bucket owner assignment for FSDP/HSDP lossless distributor (ROUND_ROBIN) — Balances per-owner-rank load for the FSDP/HSDP lossless distributor under
- [shampoo][v3+dev] Checkpoint step/lr_sum/train_mode per-param (PyTorch-conform, num_sub_groups-portable) — Store the per-group `step` / `lr_sum` / `train_mode` scalars under EVERY parameter's state
- [Shampoo][DEV][V3] Fix lossless distributor checkpoint-resume crash — The FSDP2/HSDP2 lossless distributors crash on the first optimizer step after a checkpoint resume when the model has a parameter that receives no gradient that step (a dead/frozen layer, or a conditionally-active head that is off for the batch).
- [shampoo] Drop the experimental underscore prefix and note from the now-productionized FullyShard/HybridShard lossless distributor modules (rename only, no behavior change).
- [shampoo] Balance per-owner-rank byte load in the FSDP/HSDP lossless ROUND_ROBIN distributor and correctly support `num_sub_groups > 1` via per-bucket LPT owner assignment.
- [shampoo] Store the per-group `step`, `lr_sum`, and `train_mode` scalars under every parameter's state so the checkpoint is param-keyed and DCP resume survives `num_sub_groups` changes.
- [Shampoo] Fix an FSDP2/HSDP2 lossless distributor crash on the first step after checkpoint resume when a parameter receives no gradient, by rebuilding masked lists consistently.

## 2026-07-20

- [Shampoo] Fix the DDP CIFAR-10 example gloo timeout by downloading the dataset on rank 0 only with a barrier, avoiding concurrent-pull throttling and rank desync.
- [shampoo] Wire the configured `load_balancing_config` through the HSDPDistributor's replicate-dimension block assignment instead of hardcoding the default (behavior-preserving by default).
- [shampoo] Wire the configured `load_balancing_config` through the HybridShardDistributor's replicate-dimension block assignment instead of hardcoding the default (behavior-preserving by default).

## 2026-07-16

- [shampoo][OSS] Fix two OSS CI failures by unconditionally moving test data onto the validation device and merging the license header into the existing module docstring to avoid duplicate docstrings.
- [shampoo] Fix pre-existing test-quality issues in the utils tests by hoisting in-method imports and splitting a CUDA-gated assertion so it skips rather than silently passes on CPU-only hosts.

## 2026-07-09

- [shampoo] Update the weight decay documentation to reflect the `weight_decay_type` enum instead of the old boolean `use_decoupled_weight_decay` flag.

## 2026-07-07

- [shampoo] Stabilize flaky lossless distributor/utils GPU tests by forcing NCCL onto the robust socket transport (disabling NVLS/P2P fast paths) to avoid multi-tenant contention flakes.

## 2026-07-04

- [shampoo] Extract distributor-class selection into a `_select_distributor_class` seam so a wrapper layer can override distributor selection without reimplementing instantiation (no-op refactor).

## 2026-06-26

- [Shampoo] Expose a configurable `trace_scaling_exponent` on `BaseShampooPreconditionerConfig` instead of only the boolean `use_trace_scaling` with a hardcoded square-root exponent (default preserves behavior).

## 2026-05-28

- [shampoo] Bundle small behavior-preserving cleanups in examples and test utils, including enabling DataLoader `pin_memory` and constructing the MSE target directly on the target device.
- Fix cross-rank weight decay divergence in DDP Shampoo by replacing per-rank multiplicative shrinkage with a formulation that flows decay through the all-gather so all ranks apply it uniformly.

## 2026-05-27

- [distributed_shampoo] Update the README documentation for `ClassicMomentumConfig` and the contributors list.
- [shampoo] Rebalance the HSDP param-group test fixture so greedy LPT bin-packing distributes evenly and no longer trips the per-bin `shard_size` guard.

## 2026-05-26

- [Shampoo] Fix the recompilation test to expect 2 torch.compile specializations instead of 3, since `use_grafting_method` is consumed inside a compiler-disabled region and does not cross the compile boundary.

## 2026-05-21

- [Shampoo] Remove a per-block per-step host-device sync in the `newton_schulz` normalization by replacing `max(norm, 1e-8)` with `torch.clamp`.
- [Shampoo] Remove three host-device syncs in the higher-order coupled Newton iteration (`CoupledHigherOrderConfig`) by keeping coefficients Python-side and using `torch.clamp` for the epsilon floor.

## 2026-05-15

- [shampoo] Skip `_bias_correction2` when `drop_weighting_factor_on_gsquare` is set so the dropped weighting factor is actually enforced, matching the intended prod behavior.

## 2026-05-12

- [Shampoo] Gate hot-path debug f-string interpolations behind a `logger.isEnabledFor(DEBUG)` check to avoid per-iteration host-device syncs when DEBUG is disabled.
- [Shampoo] Remove three host-device syncs by keeping intermediates as 0-d tensors instead of calling `.item`, benefiting weight-decay paths with `weight_decay > 0`.

## 2026-05-09

- [Shampoo] Extract hybrid-shard process-group creation into a shared `create_hybrid_shard_process_groups` utility used by both the HybridShard and HybridShardLossless distributors.
- [shampoo] Replace contiguous chunking in `split_param_groups` with greedy LPT bin-packing by parameter numel to balance communication volume across sub-groups.
- [shampoo] Add `num_sub_groups` to the FSDP/HSDP lossless distributor config so a user param group can be split into independent sub-groups whose collectives can overlap on-device.
- [Shampoo] Add per-param-group CUDA stream isolation (via a new `concurrency_config`) so independent NCCL communicators across param groups can overlap, with streams synced back to default after each step.

## 2026-05-04

- [shampoo][iterate averaging] Add `ClassicMomentumConfig` as an iterate-averaging option providing SGD-style momentum/dampening/Nesterov without requiring train/eval mode switching, easing migration off deprecated momentum params.
- [Shampoo] Cache `triu_indices` process-wide with int32 keys in symmetric pack/unpack to recover the ~7% QPS regression while keeping the symmetric-packing memory savings.
- [Shampoo] Add a `use_symmetric_packing` flag to store symmetric Kronecker factor matrices in flat upper-triangular form, saving roughly 33-50% of optimizer memory.

## 2026-05-01

- [Shampoo] Add a regression test asserting the torch.compile-wrapped per-group step only recompiles at the three predictable steps across all distributor variants.

## 2026-04-29

- [Shampoo] Group the parallel per-entry lists in the redistribute/gather contexts into small private Plan dataclasses to reduce misalignment while preserving the `_foreach_copy_` hot path.
- [Shampoo] Extract a shared `AllToAllContext` base class to unify the redistribute and gather contexts' send logic, eliminating roughly 60% of duplicated code.
- [Shampoo] Move fully-shard-specific code from shampoo_utils into a new shampoo_fully_shard_utils module, mirroring the existing FSDP utils pattern (pure mechanical move).

## 2026-04-27

- [Shampoo] Tensorize the iterate-averaging interpolation coefficients so Dynamo treats them as dynamic inputs, preventing per-step recompilation (and a multi-group deadlock) with ScheduleFreeConfig.

## 2026-04-15

- [optimizers] Simplify the GPA-AdamW README.
- [optimizers] Unify weight decay across GPA-AdamW and Distributed Shampoo to use multiplicative shrinkage instead of additive weight decay on the search direction.

## 2026-04-12

- [Shampoo][test] Shard training data across ranks in the tests so each rank computes a distinct gradient, making the DDP-vs-HybridShard all-reduce comparison meaningful.

## 2026-04-10

- [shampoo] Speed up the copy-in/copy-out redistribute/gather contexts with persistent send buffers and batched `foreach_copy_`, and fix an aliasing bug by cloning unpacked results.
- [Shampoo] Optimize the HybridShardLosslessDistributor by replacing per-param all-gathers and multi-call redistribution with coalesced all_to_all via the shared gather/redistribute contexts.
- [Shampoo] Introduce `GatherGradientsContext` to replace per-parameter all-gathers with a single coalesced all_to_all under ROUND_ROBIN, reducing peak gather memory to O(total_params / world_size).
- [Shampoo] Introduce `RedistributeParamsContext` to coalesce fully-shard lossless parameter redistribution into a single all_to_all with precomputed metadata and a pre-allocated recv buffer.

## 2026-04-09

- [shampoo] Port `use_trace_scaling` to dev Shampoo, normalizing factor matrices by 1/sqrt(trace) before the inverse root for better numerical stability (using `torch.where` to avoid compile graph breaks).

## 2026-04-03

- [shampoo] Reduce per-step overhead in `step` by caching `step.item` once and reusing persistent pinned/GPU learning-rate tensors instead of allocating them each step.
- [Shampoo] Document why the learning rate must be adjusted when migrating from heavy-ball momentum to GPA, including a concrete conversion example.
- Handle DTensor params/grads in DDPDistributor by calling `.to_local`, restoring compatibility with torchtitan's new `replicate` (FSDP-based) DDP path.

## 2026-04-02

- [Shampoo] Hoist `drop_weighting_factor_on_gsquare` from two subclasses up to `BaseShampooPreconditionerConfig` to eliminate duplication (leaving the separate grafting config tree unchanged).
- [Shampoo] Rename the top two PreconditionerConfig classes to `BaseShampooPreconditionerConfig` and `ClassicShampooPreconditionerConfig` to match the parallel PreconditionerList hierarchy.
- [Shampoo] Minor cleanup: remove a dead constant, drop redundant `torch.no_grad` inside already-decorated methods, and remove a redundant `tuple` wrap.
- [Shampoo] Fix misleading comments, docstrings, and typos across the distributor files (wrong class name, stale "DDP-specific" notes, FSDP/HSDP mixups).
- [Shampoo] Use a module-level `logger` in matrix_functions instead of the root `logging` functions to allow module-specific log configuration.
- [Shampoo] Add explicit parentheses to the `perform_amortized_computation` expression to make the `(A and B) or C` precedence clear.
- [Shampoo] Remove a redundant `_foreach_add_` in HybridShardDistributor's `communicate_params=True` path whose result was immediately overwritten, aligning it with the other distributors.
- [Shampoo] Fix truncated assert messages in several distributors where a misplaced closing parenthesis turned the diagnostic f-string into a dangling no-op.

## 2026-04-01

- [shampoo] Code-quality cleanup: unconditionally restore the tf32 flag, replace numpy.cumsum with itertools.accumulate to drop the numpy dependency, and add a tf32 thread-safety TODO.

## 2026-03-25

- [Shampoo][Easy] Fix minor docstring and comment typos and inaccuracies across distributed_shampoo, commons, and shampoo_utils.

## 2026-03-19

- [Shampoo][super easy] Fix minor typos and grammar in the test utils comments and docstrings.

## 2026-03-17

- [Shampoo] Rename `shampoo_checkpoint_utils` to `shampoo_state_dict_utils` and remove the unused `flatten`/`unflatten` utilities and their tests.

## 2026-03-03

- [shampoo] Add a `drop_weighting_factor_on_gsquare` option that changes the second-moment EMA update to omit the `(1 - beta2)` weighting factor on the squared gradient.
- [shampoo] Add irisz to the developers list.
- [shampoo] Add checkpoint-conversion information to the README.

## 2026-02-25

- [shampoo] Disable NCCL NVLS transport in the utils test to avoid CUDA errors during all_to_all on hardware where NVLS is unsupported.

## 2026-02-23

- [optimizers] Update the GPA README to remove internal `hpc.optimizers` references.

## 2026-02-22

- [oss][optimizers] Remove the Meta-internal Buck test command from the GPA equivalence test docstring, leaving only the open-source unittest command.

## 2026-02-21

- [oss][optimizers] Add an explicit `IterateAveragingType` enum (GPA, SCHEDULE_FREE) to cleanly separate the two modes instead of an implicit coefficient toggle, plus added equivalence tests and hyperparameter checks.

## 2026-02-20

- [oss][optimizers] Set the default GPA train and eval coefficient values to those recommended in the GPA paper.
- [oss][optimizer] Add end-to-end tests demonstrating GPA-AdamW integration with DDP and FSDP distributed training.
- [oss][optimizers] Add equivalence tests verifying GPA-AdamW matches base AdamW, LaProp, and Nesterov momentum under specific parameter settings.
- [oss][optimizers] Add GPU-based unit tests for GPA-AdamW covering numerical correctness and convergence on CUDA.
- [oss][optimizers] Add README documentation for the GPA-AdamW optimizer, including usage examples, API reference, and configuration guidelines.

## 2026-02-18

- [oss][optimizers] Add CPU unit tests for GPA-AdamW covering initialization, state dict, step computation, train/eval toggling, and edge cases, plus shared test utilities.
- [optimizers] Consolidate the six per-parallelism CIFAR-10 example scripts into a single Hydra-configured entry point using a Strategy pattern for parallelism modes.

## 2026-02-13

- [optimizers][reland] Add tests validating GPA-AdamW and Schedule-Free against a reference implementation of GPA and Schedule-Free.
- [optimizers][reland] Enforce that checkpointing does not change the train/eval mode in Shampoo.
- [optimizers][reland] Enable GPA and Schedule-Free in the CIFAR-10 examples, removing the momentum/dampening/Nesterov argparse options and toggling train/eval modes in the training loop.
- [optimizers][reland] Replace deprecated momentum/dampening in dev Shampoo with GPA and Schedule-Free iterate-averaging configs, adding a `use_masked_tensors` distributor flag to support train/eval modes.

## 2026-02-11

- Revert enabling GPA and Schedule-Free (removing momentum/dampening) in Shampoo.
- Revert enabling GPA and Schedule-Free in the CIFAR-10 examples.
- Revert enforcing that checkpointing does not change train/eval mode in Shampoo.
- Revert adding the GPA-AdamW and Schedule-Free tests.

## 2026-02-10

- [optimizers] Add tests validating GPA-AdamW and Schedule-Free against a reference implementation of GPA and Schedule-Free.
- [optimizers] Enforce that checkpointing does not change the train/eval mode in Shampoo.
- [optimizers] Enable GPA and Schedule-Free in the CIFAR-10 examples, removing the momentum/dampening/Nesterov argparse options and toggling train/eval modes in the training loop.
- [optimizers] Replace deprecated momentum/dampening in dev Shampoo with GPA and Schedule-Free iterate-averaging configs, adding a `use_masked_tensors` distributor flag to support train/eval modes.
- [shampoo] Migrate the DDP, HSDP, and Hybrid Shard distributors to the new `device_mesh._get_all_submeshes` instance-method API ahead of the old function's deprecation.

## 2026-02-09

- [shampoo] Remove the custom `distributed_state_dict`/`load_distributed_state_dict` methods now that the standard `state_dict`/`load_state_dict` work correctly.
- [shampoo] Filter out empty OptimizerModule dictionary entries (e.g. an empty `factor_matrix_indices`) from the state_dict when `store_non_tensors=False`.
- [shampoo] Increase the process-group timeout to 15 minutes for the Shampoo checkpoint tests to reduce flakiness.

## 2026-02-05

- [shampoo] Increase the distributed process-group timeout for the Shampoo checkpoint tests to reduce flakiness from remote filesystem I/O and DCP collective barriers.
- [shampoo] Add a reusable `shampoo_comm_profiler` context manager for profiling communication operations in the Shampoo distributors.

## 2026-02-02

- [shampoo] Fix a typo in dev Shampoo.

## 2026-01-26

- [shampoo][ckpt] Fix an FSDP2/HSDP2 lossless distributor checkpoint-resume mismatch by refreshing the cached full-parameter tensors that otherwise stayed stale after load.

## 2026-01-16

- [fsdp][optim] Add an `eager_nan_check` debug flag (off by default) to skip the per-step gradient NaN check, cutting Shampoo step time by ~40% and improving QPS by ~20%.

## 2026-01-14

- [shampoo] Make `rank=None` the default in local block-info construction so rank is omitted from block IDs unless a shard-dependent distributor needs it, avoiding redundant saved tensors.

## 2026-01-05

- [optimizers] Fix GPA-AdamW updating the weight sum multiple times per step by moving the update outside the per-parameter loop, and simplify the implementation for production.

## 2025-12-31

- [optimizers] Add the Generalized Primal Averaging (GPA) AdamW optimizer, an AdamW wrapper maintaining a weighted parameter average for faster, more stable convergence.

## 2025-12-19

- [optimizers] Add an independent weight decay option implementing `param = (1 - (lr / peak_lr) * weight_decay) * param` in Shampoo.
- [optimizers] Separate the weight decay tests from the Shampoo grafting tests into their own file.
- [optimizers] Add AdamC-style weight decay support to the Shampoo CIFAR-10 examples.
- [optimizers] Enable AdamC-style weight decay in Distributed Shampoo, scaling weight decay by the learning-rate-to-peak-learning-rate ratio via a new `weight_decay_type` enum and `peak_lr` field.
- [optimizers] Enable basic logging in the Shampoo CIFAR-10 examples.

## 2025-12-15

- Sync KL-Shampoo to OSS.

## 2025-12-12

- Add a simplified KL-Shampoo that adjusts only the factor-matrix update (without the per-factor eigenvalue correction), storing inverse roots or factor eigendecompositions.

## 2025-11-21

- [fsdp][optim] Add unit tests for empty local blocked-parameter lists that uncover and fix two collective-timeout issues in HSDP lossless Shampoo under ROUND_ROBIN.

## 2025-11-18

- [fsdp][optim] Disable gradient tracking for full params in hybrid-shard lossless Shampoo so `update_params` no longer fails on `distribute_tensor` for non-leaf tensors, plus a test shape fix.
- [fsdp][optim] Allow empty per-rank parameter assignments (while still rejecting empty global lists) so small models can run under ROUND_ROBIN lossless Shampoo across DDP/fully-sharded/hybrid-sharded.

## 2025-11-17

- [shampoo] Add logging of the number of all_to_all runs for lossless Shampoo ROUND_ROBIN assignment.
- [shampoo] Add an all_gather profiler to the HybridShardShampooDistributor to distinguish the DDP and HSDP all_gather paths.

## 2025-11-14

- [Shampoo] Implement a lossless HSDP distributor with round-robin parameter assignment.

## 2025-11-10

- [shampoo] Add a comment in shampoo_utils giving a padding example.

## 2025-11-08

- [shampoo][test] Add a test validating that `_local_masked_blocked_params` total elements match the original tensor under REPLICATE assignment in FullyShardLosslessDistributor.

## 2025-11-06

- [Shampoo] Implement a lossless FSDP distributor with round-robin parameter assignment, saving memory by storing only a subset of parameters per rank.

## 2025-11-05

- [Shampoo] Add an initial lossless HSDP Shampoo distributor with replicated parameters.

## 2025-11-04

- [Shampoo][OSS] Replace remaining lambda functions in the types modules with private static methods so Shampoo's state_dict param_groups can be pickled by `torch.save`.

## 2025-10-28

- [dcp][state_dict] Make the optim state-dict flatten/unflatten helpers handle arbitrarily nested dictionaries via recursion.

## 2025-10-20

- [optimizers] Revert the context manager around Newton-Schulz, porting the fix into the optimizers library.

## 2025-10-15

- [shampoo] Rename the fully-shard lossless distributor module to add a leading underscore, marking it experimental and not yet exposed.

## 2025-10-10

- [Shampoo] Update the state_dict tests to use Shampoo's new state_dict API, remove tests tied to the old `load_distributed_state_dict`, and make param_groups a list of dicts per PyTorch convention.
Loading