Algorithmically Faithful and System-Efficient Optimization for Large Language Models

Overview

Large language model training systems are highly optimized for first-order, element-wise optimizers such as AdamW. This design point has enabled large models to scale, but it also embeds a narrow execution contract: gradients are materialized by backpropagation, batch gradients may be aggregated before the optimizer step, optimizer states can be flattened and sharded, and each shard can often be updated independently. This dissertation argues that a growing class of advanced training methods breaks that contract. Zeroth-order fine-tuning estimates update directions through forward-only loss evaluations; differentially private training requires per-sample clipping before aggregation and calibrated noise on the aggregate; and matrix-based optimizers such as Muon, Shampoo, and SOAP require holistic matrix or tensor updates that conflict with distributed sharding.

The thesis develops an optimizer-aware systems view in which these methods are treated as structured computations with algorithmic invariants, memory behavior, communication objects, and scheduling constraints. Four systems instantiate this view. ZO2 makes full-parameter zeroth-order fine-tuning memory-feasible by combining CPU–GPU offloading with perturbation-consistent execution, dynamic scheduling, reusable memory, and AMP-aware transfer. DistZO2 extends ZO2 to distributed execution by introducing perturbation parallelism, scalar-gradient data parallelism, two-dimensional parallelism, and topology-aware communication. FlashDP preserves per-layer DP-SGD semantics while reformulating per-sample gradient processing as a fused, cache-friendly GPU dataflow. Canzona generalizes the thesis to distributed matrix-based optimizers by decoupling logical optimizer ownership from physical parameter distribution and by scheduling atomic optimizer tasks with load balance and asynchronous execution.

Across these systems, the dissertation shows that advanced LLM optimizers can be made practical at scale when systems preserve algorithmic faithfulness while reorganizing memory placement, communication, scheduling, and parallel execution. The broader contribution is a framework for designing LLM training systems in which optimizer semantics are visible to the runtime rather than hidden behind a fixed AdamW-like training loop.