Repository navigation
Conversation
…llbacks, speculate with clipping - Save the worker process's Adam state in checkpoints, in the ZeRO-3 CPU offload layout, and give a loaded state to the worker. - Step a sub-group that has a param without a gradient (in step()), give such a param a zero gradient, keep the gradients of earlier micro-steps, and add up a second gradient of a param in the same step. - Do not step during backward on inf/nan gradients; clamp exp_avg_sq at 0 when a step is rolled back. - Pass adam_w_mode, bias_correction and fp32_optimizer_states to the worker, allocate its gradient buffer in the master weights dtype, and step during backward only in the settings adam_rollback supports. - Step during backward also with gradient clipping, and roll back and redo the steps when clipping triggers. - Unscale fp16 gradients in the worker for steps taken during backward. - Partial offload (ZeRO-3 and SuperOffload): the backup optimizer follows the LR schedule, and its state is checkpointed. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
… commit decoupled saves - Send lr, betas, eps and weight_decay with every CPU step, so that momentum and weight decay schedules and the values of a resumed checkpoint reach the worker; a rollback uses those of the step it undoes. - Undo and redo the steps taken during backward when dynamic loss scaling raises the scale in step(). - With the decoupled checkpoint engine, commit a SuperOffload save before save_checkpoint returns: the next backward updates the fp32 weights and Adam state that the writer reads from shared memory. - Warn at initialization when /dev/shm is too small for the tensors shared with the worker. - Keep a nan exp_avg_sq nan in the rollback clamp. - Move the partial offload backup optimizer changes to their own PR. - Tests: two param groups, bf16_optimizer_states, hyperparameter changes, fp16 loss scale increase, CPU steps submitted during backward; checkpoint tests with one and two param groups. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
… sub-groups and checkpoint their Adam state With offload_optimizer.ratio < 1, the sub-groups kept on the device are stepped by a separate torch optimizer. - Give each of its param groups its own dict and copy lr, betas, eps and weight_decay from self.optimizer's group before each device step (ZeRO-3 and SuperOffload). LR schedules ignored the first group, and the shared dicts of the other groups kept every CPU partition after a checkpoint load, so device steps also updated CPU sub-groups. - Save the device sub-groups' Adam state with the other sub-groups' in optimizer_state_dict['state'], and move it to the backup optimizer on load and on universal load, so that checkpoints do not depend on the ratio and ds_to_universal finds every sub-group. - Read the state of a device sub-group from the backup optimizer in safe_get/set_full_optimizer_state. - Use torch Adam (L2) when adam_w_mode is off, and reject bias_correction=False. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
…ram dtype
torch Adam keeps the moments of the device sub-groups in the param dtype and fails on fp32 moments for bf16 params
("expected dtype float for `end`"). DeepSpeedCPUAdam keeps fp32 moments for bf16 params (fp32_optimizer_states), so a
load that keeps them in fp32 would pass fp32 moments to the backup optimizer. Cast them to the param dtype when they
move there. Test: resume with bf16 master weights at ratio 0.5.
Signed-off-by: Xinyu Lian <lian7@illinois.edu>
The rollback recovers the old moments by dividing by beta1 and beta2, so with a beta of 0 (betas=(0, 0.99), for example) undoing a CPU step after clipping, an overflow or a loss scale change wrote nan into the moments and then into the weights. A sub-group now steps during backward only when both betas of its param group are nonzero at that step; otherwise step() runs it after the checks, as for bf16 master weights. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
…ients adam_update and adam_rollback held the GIL for the whole step, so a step in one thread of a process stopped its other Python threads. They now release it, as the ZenFlow entry points already do. adam_update read the gradient in the dtype of the params: it chose the kernel by the param and state dtypes and cast the gradient pointer to the param type, so a bf16 gradient of fp32 params was read as fp32 bytes. The kernels now take the gradient type as a template parameter and are chosen by all three dtypes, which registers fp32 params and states on bf16 gradients next to the existing combinations, whose gradients keep the param dtype. - AVX-512 and NEON load the bf16 gradient as they load bf16 params and promote it in registers; AVX2, SVE and RVV builds step it in the scalar loop. - A combination without a kernel raises instead of being misread. - ZenFlow looks its kernel up by the gradient dtype too, and offsets the gradient by its own element size. Tests: another Python thread keeps waking up while a step runs, which it cannot while the step holds the GIL; fp32 params on bf16 gradients step bit for bit as on the same gradients in fp32 over 3 steps, with a scalar tail; fp16 gradients of fp32 params are rejected. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
…gradients to the CPU once SuperOffload stepped the CPU sub-groups in a worker process per rank. The process needed the fp32 weights and gradients of those sub-groups in /dev/shm (223 GiB for OPT-30B on 4 ranks), a message protocol for steps, results and checkpoints, and polling with timeouts. The gradients went into the reduce buffer, back to the GPU for the norm and out again, on the compute stream. CPU Adam: - SuperOffloadCPUAdam runs DeepSpeed's CPU Adam kernels in one thread bound to the CPU Adam cores. They release the GIL (previous commit), so the training thread keeps queueing backward kernels while a sub-group steps. The worker process, its protocol and /dev/shm go. - The Adam state stays where ZeRO-3 CPU offload keeps it, in the optimizer state of each sub-group's flat fp32 weights, so checkpoints save and load it without an override. - A step submitted during backward is a future. The CPU Adam thread waits for the sub-group's gradient copies itself, and the training thread collects the steps in the order the thread runs them; an exception in a step reaches the training thread through its future. - AMSGrad is rejected when SuperOffload starts, instead of failing in the worker at the first step. Gradients and weights: - The gradient copies to the CPU, and the finiteness flags behind them, run on a dedicated stream instead of the compute stream. - With one micro-step per step and a non-fp16 accumulation dtype, a CPU sub-group's gradient goes to the CPU once: in the master dtype, or in bf16 for bf16 gradients of fp32 master weights, which CPU Adam promotes in its kernel (previous commit). A param's second gradient in a step adds to the first in the accumulation dtype, as the reduce buffer adds it. fp16 and gradient accumulation keep the reduce buffer. - The CPU sub-groups' fp32 weights and gradients and the bf16 gradients are page-locked, registered in place where ZeRO-3 did not pin them, so the copies are asynchronous. - The weights go to the GPU in fp32 and are converted there. Those of sub-groups stepped during backward wait in a GPU staging buffer per sub-group, which step() copies into the params after the overflow and clipping checks, without unscaling and clipping those sub-groups again. - CPU Adam's scratch gradient, for fp16 loss scaling and for the rollback of a step on bf16 gradients (adam_rollback reads fp32 only), is allocated when a step needs it and is not pinned; it was 1.6 GB of pinned memory per rank at sub_group_size 4e8. Tests: the CPU buffers are page-locked; a step on a bf16 gradient matches DeepSpeedCPUAdam on the same gradient in fp32, and its rollback undoes it; a param whose gradient hook fires twice in a step, through a reentrant checkpoint, matches autograd in bf16 and fp32; the comparisons against ZeRO-3 run the bf16 path. OPT-30B, 4xB200, sub_group_size 4e8, 10 CPU Adam cores per rank, against the SuperOffload of deepspeedai#8972, both measured on master 8f9db2c: 13.82 -> 5.87 s per step; /dev/shm 223 GiB -> 0; peak host memory of the socket, /dev/shm included, 804 -> 740 GiB; GPU memory 58.6 -> 67.3 GiB per GPU, for the conversion and the staging buffers. On b62afa8 the step takes 5.86 s, and the socket's host memory peaks at 775 GiB: deepspeedai#8840 pins the master weights there in power-of-two blocks. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
The copy engine keeps serving a stream while it has a copy ready, so a copy to the CPU on another stream waited behind all the gradient copies queued on the gradient copy stream before it. Host syncs in backward make such copies: transformers' flash-attention unpadding runs nonzero() and .item() on the attention mask of a padded batch in every layer, and again in each layer recomputed in backward. With OPT-30B at batch 88 on 4xB200 (PCIe Gen5), with gradient partitions of up to 103 MB, 23 of the 96 syncs of a backward waited behind gradient copies: 463 ms per step, up to 28 ms each. A gradient larger than 16 MiB now goes in 16 MiB pieces that take turns on the copy stream and a second stream, each piece waiting for the one before. While a piece runs, neither stream has a copy ready, so the copies of other streams go in between. The copy stream ends waiting for the last piece, so the CPU steps and step(), which wait for the copy stream, still wait for the whole gradient. One B200 alone, 4 x 103 MB of bf16 gradients queued, then an .item(): - whole copies: the .item() waits 5.1 ms; the copies run at 56.3 GB/s. - 16 MiB pieces: 0.21 ms; 55.1 GB/s. 8 MiB pieces copy at 53.9 GB/s, 4 MiB at 51.7, 1 MiB at 41.8. 16 MiB pieces all on one stream, the .item() on a high-priority stream, and CUDA_DEVICE_MAX_CONNECTIONS=32 left the wait as it was. Tests: a 4090 x 4090 weight, whose bf16 gradient goes to the CPU in two pieces, matches ZeRO-3 CPU offload exactly with pinned memory. The test fails when the last piece is dropped, and when the copy stream does not wait for it. OPT-30B, 4xB200, median s per step over steps 3-5 of one run (batch 88: steps 3-4 of two runs, profiled with nsys), overlap_comm and sub_group_size 1e8 unless noted: - batch 88: 8.62 -> 8.34; one run of the code as committed: 8.34 - batch 88, unpadding computed once per batch (a benchmark patch): 8.22 -> 8.21 - batch 64, unpadding computed once per batch: 6.66 -> 6.63 - batch 32, overlap_comm off, sub_group_size 4e8, as in the previous commit: 5.83 -> 5.83 Peak GPU memory is unchanged: 90.4 GiB allocated at batch 88. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
xylian86
changed the title
…other gradient scale in one pass adam_update read the gradient as the caller left it, so a caller that unscales and clips had to write a scaled copy first: for a bf16 gradient of fp32 params, an fp32 copy and a multiply over it, a pass over the gradient before the step's own. adam_update now takes grad_scale (default 1) and multiplies each gradient element by it as it loads it, after the bf16 promotion; fp32 times fp32 rounds as the separate multiply did. AVX-512, AVX2 and NEON scale in registers; SVE and RVV builds leave scaled steps to the scalar loop. adam_rescale_step(..., old_scale, new_scale) is new. Given the weights and moments after a step on grad * old_scale, it writes those of the step on grad * new_scale from the same starting point, in one pass, with the same step count and betas: - m += (1 - beta1) * (new_scale - old_scale) * grad - v = max(0, v + (1 - beta2) * (new_scale^2 - old_scale^2) * grad^2) - p += step_size * (update(m, v) - update(m_old, v_old)); AdamW's decay of p does not depend on the gradient, so it cancels. It supports AdamW, and Adam without weight decay; L2 weight decay adds the old weights to the gradient, so it raises. Undoing a step with adam_rollback and stepping again takes two passes, and adam_rollback is a scalar loop bound by compute. 4 ranks on one socket, 10 cores each, 1e8 fp32 params with bf16 gradients: adam_rescale_step 0.051 s, as a step (0.051 s); adam_rollback 0.070 s, plus 0.040 s for its fp32 copy of the gradient. simd.h gains SIMD_MAX; max(0, v) keeps a NaN in v on every ISA. The rescaled state is the directly stepped one up to rounding: with warm moments the weights differ by 7.5e-6 of the update and the moments by about 1e-7 of their largest value. On the first step, where v holds only the step's own grad^2, the rounding of v grows like 2^-24 / c^2 for c = new_scale / old_scale (4.4e-4 of v at c = 0.01), as it does when the step is undone and taken again. Tests: grad_scale steps bf16 and fp32 gradients bit for bit as on the gradient scaled before, over 3 steps with a scalar tail; a rescale matches the direct step on the new scale, on the first and on a later step, for bf16 and fp32 gradients, and a rescale to the same scale keeps the moments bit for bit; L2 weight decay is rejected. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
…ead of undoing and redoing them SuperOffload steps the CPU sub-groups during backward, on gradients without clipping and with the loss scale of the backward. When step() then clipped, or the loss scale had changed, it undid every one of those steps with adam_rollback and stepped them again on a scaled fp32 copy of the gradient, which the training thread made on its 2 cores. Fine-tuning clips nearly every step: OPT-30B on Alpaca at batch 32 has grad norms of 3 to 15, so with gradient_clipping 1.0 every step clipped and took 20.73 s, against 5.86 s without clipping and 11.36 s for ZeRO-3 offload with the same clipping. step() now computes the factor that unscale_and_clip_grads multiplies the gradients by, with _unscale_and_clip_factor, which ZeRO-3's unscale_and_clip_grads now also uses (the same expression), and: - changes each CPU step taken during backward to that factor with one adam_rescale_step (previous commit), or leaves it when the step used that factor already; - steps the other CPU sub-groups on the gradient the worker reads, bf16 or fp32, with the factor applied as the kernel reads it, so the training thread no longer copies and scales gradients. adam_rollback stays for a skipped step (overflow or an invalid norm), for a param's second gradient in backward after its sub-group stepped, and for hyperparameters changed between backward and step(), where the step is undone and taken again. check_clip_grads, _loss_scale_changed and _cpu_step_grad go. Tests: a clipped step makes one rescale per CPU sub-group and no rollback; clipping with bf16 gradients matches ZeRO-3 CPU offload within 1e-5 over 2 steps (over 8 steps, bf16 weights rounding the other way grow the difference to 6e-5, and to 3e-5 when the steps are undone and redone); clipping in fp32 stays within 1e-5 of ZeRO-3 over 8 steps. OPT-30B, 4xB200, gradient_clipping 1.0 (every step clips), median s per step over steps 3-5: - batch 32, sub_group_size 4e8: 20.73 -> 10.34. step() took 16.8 s (waiting for the CPU steps of backward 1.7, rollbacks 7.9, gradient copies 2.7, steps 4.2) and takes 6.4 s (waiting 1.7, rescales 4.45). - batch 64 x seq 4096, sub_group_size 1e8, overlap_comm: 26.73 -> 16.05 (ZeRO-3 offload with the same clipping: 17.51). Without clipping nothing changes: OPT-350M, 50 steps with deterministic kernels, bit for bit the same losses; OPT-30B 5.86 and 11.55 s. ZeRO-3 offload with clipping keeps bit for bit the same losses too. With clipping, OPT-350M's 50-step losses differ from ZeRO-3 offload's by 3.7e-3 to 4.6e-3 per step on average, over two runs at each of two clips (undoing and redoing: 3.2e-3 to 3.4e-3; ZeRO-3 against itself with the clip changed by 1e-7: 2.7e-3 to 3.6e-3). Qwen2.5-7B, SFT on MetaMathQA (395K rows, one epoch, 68 of 3085 steps clip): the loss is within 4e-4 of ZeRO-Infinity's per step on average, GSM8K 82.6 vs 82.1. Signed-off-by: Xinyu Lian <lian7@illinois.edu>
tohtana
added
needs-gpu-test
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.