cpu-offloading

CPU offloading reduces GPU memory pressure during Megatron Bridge training by moving transformer activations and optimizer states to CPU memory. It provides configurable options for both layer-level activations and fractional optimizer state offloading to accommodate larger models.

Configuring optimizer state offloading in Megatron Bridge involves setting the optimizer.optimizer_cpu_offload key to true. You can control the memory-speed tradeoff by adjusting the optimizer.optimizer_offload_fraction and enabling optimizer.overlap_cpu_optimizer_d2h_h2d for transfer overlap.

Yes, you can offload a fraction of optimizer states using the HybridDeviceOptimizer. By setting the optimizer.optimizer_offload_fraction parameter to a value like 0.5, you partially offload Adam states to CPU, balancing memory savings with training performance.

Pipeline parallelism affects CPU activation offloading behavior. Megatron Bridge supports PP=1 activation offloading with per-transformer-layer control, while PP>1 configurations utilize optimizer offloading instead to manage memory across pipeline stages.

The best way to overlap GPU and CPU data transfers during offloading is to enable the optimizer.overlap_cpu_optimizer_d2h_h2d configuration key, which concurrently executes D2H and H2D transfers alongside GPU compute to preserve training throughput.