From 3846f50c29065ddbfa504878164bbc46ce0aa201 Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Sat, 31 Oct 2026 08:02:59 +0000 Subject: [PATCH 01/10] SuperOffload: checkpoint the CPU Adam state, fix unused params and rollbacks, 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 --- csrc/adam/cpu_adam_impl.cpp | 5 +- .../superoffload/superoffload_stage3.py | 252 ++++++++++------ .../superoffload/superoffload_utils.py | 85 +++++- deepspeed/runtime/zero/stage3.py | 41 +++ tests/unit/ops/adam/test_cpu_adam.py | 17 ++ tests/unit/v1/zero/test_invalid_grad_norm.py | 6 +- tests/unit/v1/zero/test_superoffload.py | 271 ++++++++++++++++++ 7 files changed, 576 insertions(+), 101 deletions(-) create mode 100644 tests/unit/v1/zero/test_superoffload.py diff --git a/csrc/adam/cpu_adam_impl.cpp b/csrc/adam/cpu_adam_impl.cpp index d5912272376d..80adbd5f84d4 100644 --- a/csrc/adam/cpu_adam_impl.cpp +++ b/csrc/adam/cpu_adam_impl.cpp @@ -318,7 +318,10 @@ void adamw_rollback_inplace(float* params, const float grad = grads[i]; momentum[i] = (momentum[i] - one_minus_beta1 * grad) / beta1; - variance[i] = (variance[i] - one_minus_beta2 * grad * grad) / beta2; + // exp_avg_sq is a sum of squares, so it is never negative. The inverse can leave a tiny + // negative rounding residual (in about half of the elements that were 0 before the step), + // and sqrt() of it in the next step makes the parameter NaN. + variance[i] = std::max(0.0f, (variance[i] - one_minus_beta2 * grad * grad) / beta2); } --step_count; diff --git a/deepspeed/runtime/superoffload/superoffload_stage3.py b/deepspeed/runtime/superoffload/superoffload_stage3.py index d6a4270f5aaa..ff767bc085bf 100644 --- a/deepspeed/runtime/superoffload/superoffload_stage3.py +++ b/deepspeed/runtime/superoffload/superoffload_stage3.py @@ -4,8 +4,12 @@ # DeepSpeed Team import time +from collections import defaultdict from typing import List +import torch + +from deepspeed.ops.adam import DeepSpeedCPUAdam from deepspeed.runtime.superoffload.superoffload_utils import SuperOffloadCPUOptimizer, TaskKeys, ResultKeys, EventTypes from deepspeed.runtime.zero.partition_parameters import Parameter, Tensor from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 @@ -24,6 +28,18 @@ def _validate_superoffload_accelerator(): f"SuperOffload only supports NVIDIA CUDA GPUs, but found accelerator '{accelerator.device_name()}'.") +def _cpu_adam_settings(optimizer): + """Return (adamw_mode, fp32_optimizer_states) for the DeepSpeedCPUAdam that runs ``optimizer`` in the worker.""" + if isinstance(optimizer, DeepSpeedCPUAdam): + return optimizer.adam_w_mode, optimizer.fp32_optimizer_states + if isinstance(optimizer, torch.optim.AdamW): + return True, True + if isinstance(optimizer, torch.optim.Adam): + return optimizer.defaults.get("decoupled_weight_decay", False), True + raise ValueError(f"SuperOffload steps the CPU sub-groups with DeepSpeedCPUAdam, so it supports only Adam and " + f"AdamW, but the optimizer is {type(optimizer).__name__}.") + + class SuperOffloadOptimizer_Stage3(DeepSpeedZeroOptimizer_Stage3): def __init__( @@ -36,10 +52,13 @@ def __init__( **kwargs, ): _validate_superoffload_accelerator() + adamw_mode, fp32_optimizer_states = _cpu_adam_settings(init_optimizer) self.sub_group_to_param_num = {} - self.sub_group_grad_partition_counts = {} - self._submitted_cpu_sub_groups = set() + # Per sub-group, the ds_ids of the params whose gradient for this step is in the fp32 gradient buffer. + self.sub_group_grads_ready = defaultdict(set) + # Sub-groups stepped during backward, mapped to the (lr, grad_scale) of that step for a rollback. + self._submitted_cpu_sub_groups = {} self.async_cpuadam_num = 0 self.max_grad_numel = 0 @@ -52,13 +71,22 @@ def __init__( "betas": pg["betas"], "eps": pg["eps"], "weight_decay": pg["weight_decay"], + "bias_correction": pg.get("bias_correction", True), "amsgrad": pg["amsgrad"], }) + # A step taken during backward is undone with adam_rollback, which supports only fp32 AdamW with bias + # correction. With other settings, step() runs the CPU sub-groups after the overflow and clipping checks. + self.speculative_cpu_step = (adamw_mode and fp32_optimizer_states + and self.master_weights_and_grads_dtype == torch.float32 + and all(config["bias_correction"] for config in optimizer_configs)) cpuadam_cores_perc = kwargs.get("cpuadam_cores_perc", 0.8) self.superoffload_cpu_optimizer = SuperOffloadCPUOptimizer(optimizer_config=optimizer_configs, cpuadam_cores_perc=cpuadam_cores_perc, max_grad_numel=self.max_grad_numel, - pin_memory=self.offload_optimizer_pin_memory) + pin_memory=self.offload_optimizer_pin_memory, + adamw_mode=adamw_mode, + fp32_optimizer_states=fp32_optimizer_states, + grad_dtype=self.master_weights_and_grads_dtype) def _create_fp16_sub_groups(self, params_group): @@ -103,15 +131,11 @@ def step_with_gradscaler(optimizer): cur_device = self.subgroup_to_device[sub_group_id] if cur_device != 'cpu': + self._sync_backup_optimizer_group(param_group_id) self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] step_with_gradscaler(self.backup_optimizer) self.backup_optimizer.param_groups[param_group_id]['params'] = [] - @instrument_w_nvtx - def independent_gradient_partition_epilogue(self): - super().independent_gradient_partition_epilogue() - self.sub_group_grad_partition_counts.clear() - @instrument_w_nvtx def _reassign_or_swap_out_partitioned_parameters(self, sub_group_id): if self.subgroup_to_device[sub_group_id] == 'cpu': @@ -142,7 +166,12 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L # Accumulate gradient into the grad_buffer, mirroring base class logic grad_buffer = self._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[param.ds_id].narrow( 0, 0, grad_partition.numel()) - if self.micro_step_id == 0: + # The buffer holds an older step's gradient until the param's first gradient of this step. Testing the + # param rather than micro_step_id also keeps both gradients of a param that two reentrant checkpoint + # segments use. + first_grad = param.ds_id not in self._offload_boundary_param_ids + self._offload_boundary_param_ids.add(param.ds_id) + if first_grad: grad_buffer.copy_(grad_partition, non_blocking=True) grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) elif get_accelerator().on_accelerator(grad_buffer): @@ -154,30 +183,39 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L grad_buffer = cuda_grad_buffer if self.is_gradient_accumulation_boundary(): + if i in self._submitted_cpu_sub_groups: + # Another gradient for a sub-group that has already stepped: undo that step while the fp32 + # buffer still holds the gradient it used. + self._wait_for_async_operations() + self._rollback_cpu_sub_group(i) self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow( 0, dest_offset, grad_buffer.numel()) fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype), non_blocking=True) - self.sub_group_grad_partition_counts[i] = self.sub_group_grad_partition_counts.get(i, 0) + 1 - if self.sub_group_grad_partition_counts[i] == self.sub_group_to_param_num[i]: - completed_sub_groups.append(i) + self.sub_group_grads_ready[i].add(param.ds_id) + if len(self.sub_group_grads_ready[i]) == self.sub_group_to_param_num[i]: + completed_sub_groups.append(i) if self.is_gradient_accumulation_boundary() and completed_sub_groups: get_accelerator().current_stream().synchronize() - for i in completed_sub_groups: - if self.subgroup_to_device[i] == 'cpu' and not self.clip_grad: + # A param that gets a second gradient adds its sub-group again. + for i in dict.fromkeys(completed_sub_groups): + if self.subgroup_to_device[i] == 'cpu' and self.speculative_cpu_step and self._grads_finite(i): param_group_id = self.sub_group_to_group_id[i] fp32_param = self.fp32_partitioned_groups_flat[i] current_lr = self.optimizer.param_groups[param_group_id]['lr'] + # The gradients still carry the loss scale, which step() would remove. + grad_scale = 1.0 / self.loss_scale self.superoffload_cpu_optimizer.async_step(param_group_id, i, fp32_param.data, fp32_param.grad.data, - lr=current_lr) - self._submitted_cpu_sub_groups.add(i) + lr=current_lr, + grad_scale=grad_scale) + self._submitted_cpu_sub_groups[i] = (current_lr, grad_scale) self.async_cpuadam_num += 1 result = self.superoffload_cpu_optimizer.get_result() @@ -198,14 +236,14 @@ def step(self, closure=None): Not supporting closure. """ self._wait_for_async_operations() + self._complete_grads() self._pre_step() self._partition_all_parameters() self._overflow_check_and_loss_scale_update(update_scale=False) if self.overflow: - if not self.clip_grad: - self._handle_overflow_rollback() + self._rollback_submitted_cpu_sub_groups() self._loss_scale_update_and_overflow_cleanup() return @@ -213,8 +251,7 @@ def step(self, closure=None): scaled_global_grad_norm = combine_grad_norm_groups(norm_groups) if is_invalid_grad_norm(scaled_global_grad_norm): self.overflow = True - if not self.clip_grad: - self._handle_overflow_rollback() + self._rollback_submitted_cpu_sub_groups() self._loss_scale_update_and_overflow_cleanup() return self._loss_scale_update_and_overflow_cleanup() @@ -224,35 +261,17 @@ def step(self, closure=None): timer_names.add(OPTIMIZER_STEP_TIMER) self.timers(OPTIMIZER_STEP_TIMER).start() - if self.clip_grad: - self._step_with_clipping(scaled_global_grad_norm, timer_names) - else: - self._step_without_clipping(scaled_global_grad_norm, timer_names) - - self.timers(OPTIMIZER_STEP_TIMER).stop() - self._post_step(timer_names) - - def _post_step(self, timer_names): - super()._post_step(timer_names) - self._submitted_cpu_sub_groups.clear() - - def _step_without_clipping(self, scaled_global_grad_norm, timer_names): - """Fast path: async CPU steps already completed during backward.""" - for sub_group_id, group in enumerate(self.fp16_groups): - self._prepare_sub_group(sub_group_id, timer_names) - self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) - self._optimizer_step(sub_group_id) - self._reassign_or_swap_out_partitioned_parameters(sub_group_id) - self._release_sub_group(sub_group_id, timer_names) + if self.check_clip_grads(scaled_global_grad_norm): + # The steps taken during backward used the unclipped gradients. + self._rollback_submitted_cpu_sub_groups() - def _step_with_clipping(self, scaled_global_grad_norm, timer_names): - """Clipping path: no async steps were done during backward, - so we unscale+clip first, then step all sub-groups.""" for sub_group_id, group in enumerate(self.fp16_groups): self._prepare_sub_group(sub_group_id, timer_names) self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) - if self.subgroup_to_device[sub_group_id] == 'cpu': + if self.subgroup_to_device[sub_group_id] != 'cpu': + self._optimizer_step(sub_group_id) + elif sub_group_id not in self._submitted_cpu_sub_groups: param_group_id = self.sub_group_to_group_id[sub_group_id] fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] current_lr = self.optimizer.param_groups[param_group_id]['lr'] @@ -261,12 +280,53 @@ def _step_with_clipping(self, scaled_global_grad_norm, timer_names): fp32_param.data, fp32_param.grad.data, lr=current_lr) - else: - self._optimizer_step(sub_group_id) self._reassign_or_swap_out_partitioned_parameters(sub_group_id) self._release_sub_group(sub_group_id, timer_names) + self.timers(OPTIMIZER_STEP_TIMER).stop() + self._post_step(timer_names) + + def _post_step(self, timer_names): + super()._post_step(timer_names) + self._submitted_cpu_sub_groups.clear() + + def _complete_grads(self): + """Fill the fp32 gradient of the params that got no gradient in the boundary micro-step. + + A param used only in earlier micro-steps has its gradient in the reduce buffer. A param not used in this step + gets a zero gradient, as set_none_gradients_to_zero gives it in ZeRO-1/2; zeroing its reduce buffer also keeps + an older inf out of the fp16 overflow check. Their sub-groups did not step during backward, so step() runs + them. + """ + filled = False + for i, group in enumerate(self.fp16_groups): + ready = self.sub_group_grads_ready[i] + if len(ready) == self.sub_group_to_param_num[i]: + continue + for param in group: + if param.ds_id in ready: + continue + _, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + grad_buffer = self._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[param.ds_id].narrow( + 0, 0, num_elements) + if param.ds_id not in self._offload_boundary_param_ids: + grad_buffer.zero_() + if not get_accelerator().on_accelerator(grad_buffer): + grad_buffer = grad_buffer.to(get_accelerator().current_device_name(), non_blocking=True) + self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) + fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow(0, dest_offset, num_elements) + fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype), non_blocking=True) + filled = True + if filled: + get_accelerator().current_stream().synchronize() + self.sub_group_grads_ready.clear() + + def _grads_finite(self, sub_group_id): + """Whether the sub-group's gradients are finite. A step on inf or nan cannot be rolled back.""" + norms = [self.norm_for_param_grads[self.get_param_id(param)] for param in self.fp16_groups[sub_group_id]] + return bool(torch.isfinite(torch.stack(norms)).all()) + def _wait_for_async_operations(self, timeout_seconds=60): """Wait for all pending asynchronous CPU optimizer operations to complete with timeout error. @@ -334,56 +394,76 @@ def _sync_cpu_optimizer_step(self, fp32_grad_data, rollback: bool = False, lr: float = None, - timeout_seconds: int = 60): + timeout_seconds: int = 60, + grad_scale: float = 1.0): event_type = EventTypes.ROLLBACK if rollback else EventTypes.ADAM_STEP self.superoffload_cpu_optimizer.async_step(param_group_id, sub_group_id, fp32_param_data, fp32_grad_data, rollback=rollback, - lr=lr) + lr=lr, + grad_scale=grad_scale) # Wait for completion self._wait_for_single_async_result(event_type, timeout_seconds) - def _handle_overflow_rollback(self): - """Handle gradient overflow by rolling back CPU optimizer states.""" - for sub_group_id in self._submitted_cpu_sub_groups: - param_group_id = self.sub_group_to_group_id[sub_group_id] - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - self._sync_cpu_optimizer_step(param_group_id, - sub_group_id, - fp32_param.data, - fp32_param.grad.data, - rollback=True) - self._submitted_cpu_sub_groups.clear() - - def _handle_gradient_clipping(self, scaled_global_grad_norm): - """Handle gradient clipping with CPU optimizer rollback and re-optimization.""" - for sub_group_id, _ in enumerate(self.fp16_groups): - if self.subgroup_to_device[sub_group_id] == 'cpu': - param_group_id = self.sub_group_to_group_id[sub_group_id] - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + def _rollback_submitted_cpu_sub_groups(self): + """Undo the CPU steps taken during backward, for a skipped step or before redoing them with clipping.""" + for sub_group_id in list(self._submitted_cpu_sub_groups): + self._rollback_cpu_sub_group(sub_group_id) - # Rollback CPU optimizer states - self._sync_cpu_optimizer_step(param_group_id, - sub_group_id, - fp32_param.data, - fp32_param.grad.data, - rollback=True) - - # Clip gradients and re-optimize - self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) - - current_lr = self.optimizer.param_groups[param_group_id]['lr'] - self._sync_cpu_optimizer_step(param_group_id, - sub_group_id, - fp32_param.data, - fp32_param.grad.data, - rollback=False, - lr=current_lr) + def _rollback_cpu_sub_group(self, sub_group_id): + lr, grad_scale = self._submitted_cpu_sub_groups.pop(sub_group_id) + param_group_id = self.sub_group_to_group_id[sub_group_id] + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + self._sync_cpu_optimizer_step(param_group_id, + sub_group_id, + fp32_param.data, + fp32_param.grad.data, + rollback=True, + lr=lr, + grad_scale=grad_scale) @instrument_w_nvtx def check_clip_grads(self, total_norm): - """Check if gradients need to be clipped based on the global norm.""" - unscaled_norm = total_norm / self.loss_scale - return self.clip_grad and unscaled_norm > self.clip_grad + """Check if unscale_and_clip_grads will scale the gradients down, with the same test.""" + if not self.clip_grad: + return False + clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad + return bool(clip > 1.0) + + def state_dict(self): + # The worker process holds the Adam state of the CPU sub-groups. Save it where ZeRO-3 CPU offload keeps it. + self._wait_for_async_operations() + cpu_params = [] + for sub_group_id, state in self.superoffload_cpu_optimizer.get_state().items(): + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + self.optimizer.state[fp32_param] = state + cpu_params.append(fp32_param) + try: + return super().state_dict() + finally: + for fp32_param in cpu_params: + del self.optimizer.state[fp32_param] + + def load_state_dict(self, + state_dict_list, + load_optimizer_states=True, + load_from_fp32_weights=False, + checkpoint_folder=None, + load_serial=None, + param_shapes=None): + super().load_state_dict(state_dict_list, + load_optimizer_states=load_optimizer_states, + load_from_fp32_weights=load_from_fp32_weights, + checkpoint_folder=checkpoint_folder, + load_serial=load_serial, + param_shapes=param_shapes) + if not load_optimizer_states: + return + # Move the loaded Adam state of the CPU sub-groups to the worker process, which steps them. + cpu_state = {} + for sub_group_id, fp32_param in enumerate(self.fp32_partitioned_groups_flat): + if self.subgroup_to_device[sub_group_id] == 'cpu' and fp32_param in self.optimizer.state: + cpu_state[sub_group_id] = self.optimizer.state.pop(fp32_param) + self.superoffload_cpu_optimizer.set_state(cpu_state) diff --git a/deepspeed/runtime/superoffload/superoffload_utils.py b/deepspeed/runtime/superoffload/superoffload_utils.py index 266f8e922a2c..83c924495063 100644 --- a/deepspeed/runtime/superoffload/superoffload_utils.py +++ b/deepspeed/runtime/superoffload/superoffload_utils.py @@ -7,6 +7,7 @@ """ +import time from typing import Dict, Optional, Any import torch import torch.multiprocessing as mp @@ -24,32 +25,44 @@ class TaskKeys: SUB_GROUP_ID = "sub_group_id" ROLLBACK = "rollback" LR = "lr" + GRAD_SCALE = "grad_scale" + COMMAND = "command" + STATE = "state" class ResultKeys: UPDATED_PARAM = "updated_param" EVENT_TYPE = "event_type" + STATE = "state" class EventTypes: ADAM_STEP = "adam_step" ROLLBACK = "rollback" + GET_STATE = "get_state" + SET_STATE = "set_state" -def _allocate_worker_grad_buffer(max_grad_numel: int, pin_memory: bool) -> torch.Tensor: +def _allocate_worker_grad_buffer(max_grad_numel: int, + pin_memory: bool, + dtype: torch.dtype = torch.float32) -> torch.Tensor: # Scratch destination for the GPU->CPU grad copy. Pinning is optional: it # speeds DMA but is not required for correctness, so honor # offload_optimizer.pin_memory rather than always page-locking. + # dtype is the master weights dtype: the optimizer requires the gradient to match the parameter. if pin_memory: - return get_accelerator().pin_empty(max_grad_numel, dtype=torch.float32, device='cpu') - return torch.empty(max_grad_numel, dtype=torch.float32, device='cpu') + return get_accelerator().pin_empty(max_grad_numel, dtype=dtype, device='cpu') + return torch.empty(max_grad_numel, dtype=dtype, device='cpu') def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, result_queue: mp.SimpleQueue, optimizer_config: Dict[str, Any], max_grad_numel: int, - pin_memory: bool = True) -> None: + pin_memory: bool = True, + adamw_mode: bool = True, + fp32_optimizer_states: bool = True, + grad_dtype: torch.dtype = torch.float32) -> None: """ This function runs in a separate process and continuously processes optimization tasks from the parameter queue. It creates a DeepSpeedCPUAdam optimizer and @@ -59,9 +72,12 @@ def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, param_queue: Queue for receiving optimization tasks result_queue: Queue for sending back optimization results optimizer_config: Configuration dictionary for the optimizer containing - lr, betas, eps, weight_decay, and amsgrad parameters + lr, betas, eps, weight_decay, bias_correction, and amsgrad parameters max_grad_numel: Maximum number of elements expected in gradient tensors pin_memory: Whether to page-lock the reusable grad scratch buffer + adamw_mode: Whether the optimizer is AdamW (decoupled weight decay) or Adam + fp32_optimizer_states: Whether to keep the optimizer states in fp32 for low-precision parameters + grad_dtype: dtype of the master weights and their gradients """ cpu_tensor = torch.randn(1, device="cpu") cpu_param = torch.nn.Parameter(cpu_tensor) @@ -75,10 +91,13 @@ def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, first_cfg = pg_configs[0] optimizer = DeepSpeedCPUAdam([cpu_param], lr=first_cfg["lr"], + bias_correction=first_cfg["bias_correction"], betas=first_cfg["betas"], eps=first_cfg["eps"], weight_decay=first_cfg["weight_decay"], - amsgrad=first_cfg["amsgrad"]) + amsgrad=first_cfg["amsgrad"], + adamw_mode=adamw_mode, + fp32_optimizer_states=fp32_optimizer_states) for cfg in pg_configs[1:]: dummy = torch.nn.Parameter(torch.randn(1, device="cpu")) optimizer.add_param_group({ @@ -87,6 +106,7 @@ def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, "betas": cfg["betas"], "eps": cfg["eps"], "weight_decay": cfg["weight_decay"], + "bias_correction": cfg["bias_correction"], "amsgrad": cfg["amsgrad"], }) except KeyError as e: @@ -96,7 +116,7 @@ def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, return # Pre-allocate reusable host buffer for gradients - grad_buffer = _allocate_worker_grad_buffer(max_grad_numel, pin_memory) + grad_buffer = _allocate_worker_grad_buffer(max_grad_numel, pin_memory, grad_dtype) while True: try: @@ -106,12 +126,26 @@ def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, logger.debug("Received termination signal, shutting down worker") break + # The Adam state lives in this process, so checkpoints get and set it here. Its tensors move to shared + # memory when they are sent, without a copy that stays alive. + command = task.get(TaskKeys.COMMAND) + if command == EventTypes.GET_STATE: + state = {sub_group_id: state for sub_group_id, state in optimizer.state.items() if state} + result_queue.put({ResultKeys.EVENT_TYPE: EventTypes.GET_STATE, ResultKeys.STATE: state}) + continue + if command == EventTypes.SET_STATE: + optimizer.state.clear() + optimizer.state.update(task[TaskKeys.STATE]) + result_queue.put({ResultKeys.EVENT_TYPE: EventTypes.SET_STATE}) + continue + param_data = task[TaskKeys.PARAM_DATA] param_grad = task[TaskKeys.PARAM_GRAD] param_group_id = task[TaskKeys.PARAM_GROUP_ID] sub_group_id = task[TaskKeys.SUB_GROUP_ID] rollback = task.get(TaskKeys.ROLLBACK, False) task_lr = task.get(TaskKeys.LR, None) + grad_scale = task.get(TaskKeys.GRAD_SCALE, 1.0) logger.debug(f"Processing param_group_id: {param_group_id}, sub_group_id: {sub_group_id}") @@ -131,7 +165,11 @@ def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, break param_grad_cpu = grad_buffer[:grad_numel].view_as(param_grad) - param_grad_cpu.copy_(param_grad, non_blocking=False) + if grad_scale == 1.0: + param_grad_cpu.copy_(param_grad, non_blocking=False) + else: + # A step submitted during backward gets the gradient before the loss scale is removed. + torch.mul(param_grad, grad_scale, out=param_grad_cpu) fp32_param = torch.nn.Parameter(param_data) fp32_param.grad = param_grad_cpu @@ -185,7 +223,10 @@ def __init__(self, optimizer_config: Dict[str, Any], cpuadam_cores_perc: float = 0.8, max_grad_numel: int = 1000000, - pin_memory: bool = True) -> None: + pin_memory: bool = True, + adamw_mode: bool = True, + fp32_optimizer_states: bool = True, + grad_dtype: torch.dtype = torch.float32) -> None: if not 0 < cpuadam_cores_perc <= 1: raise ValueError("cpuadam_cores_perc must be between 0 and 1") @@ -196,7 +237,8 @@ def __init__(self, self.cpuadam_process = self.mp_context.Process( target=superoffload_optimizer_worker, - args=(self.param_queue, self.result_queue, optimizer_config, max_grad_numel, pin_memory), + args=(self.param_queue, self.result_queue, optimizer_config, max_grad_numel, pin_memory, adamw_mode, + fp32_optimizer_states, grad_dtype), daemon=True, ) self.cpuadam_process.start() @@ -241,7 +283,8 @@ def async_step(self, fp32_param: torch.Tensor, fp32_grad: torch.Tensor, rollback: bool = False, - lr: float = None) -> None: + lr: float = None, + grad_scale: float = 1.0) -> None: """ Queue parameter for optimization in the worker process. """ @@ -257,7 +300,27 @@ def async_step(self, } if lr is not None: task[TaskKeys.LR] = lr + if grad_scale != 1.0: + task[TaskKeys.GRAD_SCALE] = grad_scale + self.param_queue.put(task) + + def _request(self, task: Dict[str, Any], event_type: str) -> Dict[str, Any]: + """Send a command to the worker process and wait for its reply. Call it with no step in flight.""" self.param_queue.put(task) + while self.result_queue.empty(): + if not self.cpuadam_process.is_alive(): + raise RuntimeError("Worker process is not alive") + time.sleep(0.001) + return self.get_result(expected_event_type=event_type) + + def get_state(self) -> Dict[int, Dict[str, Any]]: + """Return the Adam state of each sub-group that has taken a step, keyed by sub-group id.""" + result = self._request({TaskKeys.COMMAND: EventTypes.GET_STATE}, EventTypes.GET_STATE) + return result[ResultKeys.STATE] + + def set_state(self, state: Dict[int, Dict[str, Any]]) -> None: + """Replace the Adam state of the worker, for example with the state loaded from a checkpoint.""" + self._request({TaskKeys.COMMAND: EventTypes.SET_STATE, TaskKeys.STATE: state}, EventTypes.SET_STATE) def get_result(self, expected_event_type: str = None) -> Optional[Dict[str, Any]]: """ diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 27ddf8eec69b..4a1a5d699d85 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -1239,6 +1239,19 @@ def _create_fp16_sub_groups(self, params_group): return sub_groups + def _sync_backup_optimizer_group(self, param_group_id): + """Copy the hyperparameters of a param group to the optimizer of the sub-groups kept on the device. + + LR schedulers and set_lr() update self.optimizer only. The backup optimizer has its own first param group, + and after a checkpoint load its own copy of every group. + """ + backup_group = self.backup_optimizer.param_groups[param_group_id] + main_group = self.optimizer.param_groups[param_group_id] + if backup_group is main_group: + return + for key in ("lr", "betas", "eps", "weight_decay"): + backup_group[key] = main_group[key] + def _optimizer_step(self, sub_group_id): param_group_id = self.sub_group_to_group_id[sub_group_id] fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] @@ -1260,6 +1273,7 @@ def step_with_gradscaler(optimizer): step_with_gradscaler(self.optimizer) self.optimizer.param_groups[param_group_id]['params'] = [] else: + self._sync_backup_optimizer_group(param_group_id) self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] step_with_gradscaler(self.backup_optimizer) self.backup_optimizer.param_groups[param_group_id]['params'] = [] @@ -3440,6 +3454,20 @@ def _clear_fp32_optimizer_param_groups(self): for param_group in self.optimizer.param_groups: param_group['params'] = [] + def _set_backup_optimizer_param_groups(self): + # A new backup optimizer still holds its placeholder param, and after a load a group can share the param list + # of a replaced self.optimizer group. + self._clear_backup_optimizer_param_groups() + for sub_group_id, _ in enumerate(self.fp16_groups): + if self.subgroup_to_device[sub_group_id] != 'cpu': + param_group_id = self.sub_group_to_group_id[sub_group_id] + self.backup_optimizer.param_groups[param_group_id]['params'].append( + self.fp32_partitioned_groups_flat[sub_group_id]) + + def _clear_backup_optimizer_param_groups(self): + for param_group in self.backup_optimizer.param_groups: + param_group['params'] = [] + def _zero3_partition_group_metadata(self): metadata = [] has_autoep_expert_group = False @@ -3474,6 +3502,12 @@ def _rigid_state_dict(self): state_dict[FP32_FLAT_GROUPS] = self.fp32_partitioned_groups_flat self._clear_fp32_optimizer_param_groups() + # With partial offload, the sub-groups kept on the device have their Adam state in the backup optimizer. + if getattr(self, "backup_optimizer", None) is not None: + self._set_backup_optimizer_param_groups() + state_dict['backup_optimizer_state_dict'] = self.backup_optimizer.state_dict() + self._clear_backup_optimizer_param_groups() + return state_dict def state_dict(self): @@ -3583,6 +3617,13 @@ def _rigid_load_state_dict(self, state_dict, load_optimizer_states=True): self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) self._clear_fp32_optimizer_param_groups() + # Checkpoints saved without partial offload, or before its state was saved, do not have it. + has_backup_state = 'backup_optimizer_state_dict' in state_dict + if load_optimizer_states and has_backup_state and getattr(self, "backup_optimizer", None) is not None: + self._set_backup_optimizer_param_groups() + self.backup_optimizer.load_state_dict(state_dict['backup_optimizer_state_dict']) + self._clear_backup_optimizer_param_groups() + if self.swap_optimizer: # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint self.optimizer_swapper.purge_state() diff --git a/tests/unit/ops/adam/test_cpu_adam.py b/tests/unit/ops/adam/test_cpu_adam.py index 762049e0abfd..ca925c23c203 100644 --- a/tests/unit/ops/adam/test_cpu_adam.py +++ b/tests/unit/ops/adam/test_cpu_adam.py @@ -443,6 +443,23 @@ def test_rollback_subgroup_basic(self, dtype): # Verify step count decremented assert optimizer.state[subgroup_id]['step'] == 0, "Step count should be decremented after rollback" + def test_rollback_of_first_step_keeps_exp_avg_sq_non_negative(self): + """exp_avg_sq is 0 before the first step. Rolling that step back must not leave it negative, because sqrt() + of a negative value in the next step turns the parameter into NaN.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + model_size = 4096 + param = torch.nn.Parameter(torch.randn(model_size, device='cpu')) + optimizer = DeepSpeedCPUAdam([param]) + param.grad = torch.randn(model_size, device='cpu') + + optimizer.step_subgroup(0) + optimizer.rollback_subgroup(0) + assert (optimizer.state[0]['exp_avg_sq'] >= 0).all() + + optimizer.step_subgroup(0) + assert torch.isfinite(param.data).all() + def test_rollback_subgroup_uninitialized_error(self): """Test that rollback_subgroup raises error for uninitialized subgroup.""" from deepspeed.ops.adam import DeepSpeedCPUAdam diff --git a/tests/unit/v1/zero/test_invalid_grad_norm.py b/tests/unit/v1/zero/test_invalid_grad_norm.py index 2039e549d306..4d9bf0449166 100644 --- a/tests/unit/v1/zero/test_invalid_grad_norm.py +++ b/tests/unit/v1/zero/test_invalid_grad_norm.py @@ -154,18 +154,18 @@ def test_bf16_optimizer_reports_invalid_step_and_recovers(monkeypatch): def test_superoffload_rolls_back_only_subgroups_submitted_this_step(): optimizer = object.__new__(SuperOffloadOptimizer_Stage3) - optimizer._submitted_cpu_sub_groups = {1} + optimizer._submitted_cpu_sub_groups = {1: (0.1, 1.0)} optimizer.sub_group_to_group_id = {0: 10, 1: 11} parameter = SimpleNamespace(data=torch.tensor([1.0]), grad=SimpleNamespace(data=torch.tensor([2.0]))) optimizer.fp32_partitioned_groups_flat = [parameter, parameter] rollbacks = [] optimizer._sync_cpu_optimizer_step = lambda *args, **kwargs: rollbacks.append((args, kwargs)) - optimizer._handle_overflow_rollback() + optimizer._rollback_submitted_cpu_sub_groups() assert [args[1] for args, _ in rollbacks] == [1] assert rollbacks[0][1]["rollback"] - assert optimizer._submitted_cpu_sub_groups == set() + assert optimizer._submitted_cpu_sub_groups == {} class TestZeROBFloat16Stage3InvalidNormRecovery(DistributedTest): diff --git a/tests/unit/v1/zero/test_superoffload.py b/tests/unit/v1/zero/test_superoffload.py new file mode 100644 index 000000000000..06a91093d0ad --- /dev/null +++ b/tests/unit/v1/zero/test_superoffload.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""SuperOffload against ZeRO-3 CPU offload, which applies the same DeepSpeedCPUAdam update in the main process.""" + +import pytest +import torch + +import deepspeed +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import CPUAdamBuilder +from unit.common import DistributedTest + +if not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: + pytest.skip("cpu-adam is not compatible", allow_module_level=True) +if get_accelerator().device_name() != "cuda": + pytest.skip("SuperOffload supports only CUDA", allow_module_level=True) + +HIDDEN = 64 + + +class PartlyUsedModel(torch.nn.Module): + """`a` is always used. `b` and `c` are used, skipped, or multiplied by 0, which gives them a zero gradient.""" + + def __init__(self): + super().__init__() + self.a = torch.nn.Linear(HIDDEN, HIDDEN) + self.b = torch.nn.Linear(HIDDEN, HIDDEN) + self.c = torch.nn.Linear(HIDDEN, HIDDEN) + + def forward(self, x, uses): + h = torch.nn.functional.gelu(self.a(x)) + for name in ("b", "c"): + use = uses.get(name, "use") + if use == "use": + h = h + getattr(self, name)(x) + elif use == "zero": + h = h + 0 * getattr(self, name)(x) + return h.float().pow(2).mean() + + +def _config(super_offload, clip=0.0, ga=1, precision="bf16", ratio=1.0, sub_group_size=1e9, **optimizer_params): + optimizer_params.setdefault("lr", 3e-3) + optimizer_params.setdefault("weight_decay", 0.01) + optimizer_type = optimizer_params.pop("type", "AdamW") + config = { + "train_micro_batch_size_per_gpu": 4, + "gradient_accumulation_steps": ga, + "gradient_clipping": clip, + "steps_per_print": 1000, + "optimizer": { + "type": optimizer_type, + "params": optimizer_params + }, + "zero_optimization": { + "stage": 3, + "sub_group_size": sub_group_size, + "offload_optimizer": { + "device": "cpu", + "ratio": ratio, + "super_offload": super_offload, + }, + }, + } + if precision == "fp16": + config["fp16"] = {"enabled": True, "initial_scale_power": 24} + elif precision != "fp32": + config["bf16"] = {"enabled": True, "bf16_master_weights_and_grads": precision == "bf16_master"} + return config + + +def _engine(config): + torch.manual_seed(0) + model = PartlyUsedModel() + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + return engine + + +def _masters(engine): + return torch.cat([t.detach().float().flatten().cpu() for t in engine.optimizer.fp32_partitioned_groups_flat]) + + +def _close(engine): + if hasattr(engine.optimizer, "superoffload_cpu_optimizer"): + engine.optimizer.superoffload_cpu_optimizer.close() + + +def _train(engine, steps, uses_for, first_step=0): + """Train `steps` optimizer steps and return the master weights after each one. + + uses_for(step, micro_step) gives the `uses` of PartlyUsedModel.forward. + """ + dtype = torch.float32 + if engine.fp16_enabled(): + dtype = torch.float16 + elif engine.bfloat16_enabled(): + dtype = torch.bfloat16 + generator = torch.Generator().manual_seed(1 + first_step + 100 * dist.get_rank()) + masters = [] + for step in range(first_step, first_step + steps): + for micro_step in range(engine.gradient_accumulation_steps()): + x = torch.randn(4, HIDDEN, generator=generator).to(engine.device, dtype) + loss = engine(x, uses_for(step, micro_step)) + engine.backward(loss) + engine.step() + masters.append(_masters(engine)) + return masters + + +def _all_used(step, micro_step): + return {} + + +def _b_unused_on_even_steps(unused): + return lambda step, micro_step: {"b": "use" if step % 2 else unused} + + +def _b_and_c_in_one_micro_step(unused): + return lambda step, micro_step: {"b": unused, "c": "use"} if micro_step else {"b": "use", "c": unused} + + +# name: (config overrides, uses_for(unused), compare exactly) +CASES = { + "all_used": (dict(), lambda unused: _all_used, True), + "unused_param": (dict(), _b_unused_on_even_steps, True), + # Clipping is on but never triggers, so the CPU steps still run during backward and must stay exact. + "unused_param_clip": (dict(clip=100.0), _b_unused_on_even_steps, True), + "grad_accumulation": (dict(ga=2), _b_and_c_in_one_micro_step, True), + "adam": (dict(type="Adam", adam_w_mode=False), lambda unused: _all_used, True), + "no_bias_correction": (dict(bias_correction=False), lambda unused: _all_used, True), + "bf16_master_weights": (dict(precision="bf16_master"), lambda unused: _all_used, True), + # Clipping undoes the steps taken during backward and redoes them, which is exact up to rounding. In fp32 the + # rounding stays small; bf16 weights would amplify it. + "clipping": (dict(clip=0.05, precision="fp32"), lambda unused: _all_used, False), + # Overflow steps are skipped; the steps taken during backward are undone, exact up to rounding. + "fp16_overflow": (dict(precision="fp16", sub_group_size=4500), lambda unused: _all_used, False), +} + + +class TestSuperOffloadMatchesZeRO3(DistributedTest): + world_size = 1 + # SuperOffload starts its CPU worker process, which a daemonic test process cannot do. + non_daemonic_procs = True + + @pytest.mark.parametrize("case", list(CASES)) + def test(self, case): + overrides, uses_for, exact = CASES[case] + # The reference gives an unused param an exactly zero gradient, which SuperOffload must give a param that + # gets no gradient (as ZeRO-1/2 do). + engine = _engine(_config(super_offload=False, **overrides)) + expected = _train(engine, 8, uses_for("zero")) + _close(engine) + engine = _engine(_config(super_offload=True, **overrides)) + actual = _train(engine, 8, uses_for("skip")) + _close(engine) + + for step, (a, e) in enumerate(zip(actual, expected)): + assert torch.isfinite(a).all(), f"step {step}" + if exact: + assert torch.equal(a, e), f"step {step}: max diff {(a - e).abs().max().item()}" + else: + torch.testing.assert_close(a, e, rtol=0, atol=1e-5, msg=f"step {step}") + + +class InfInjectionModel(torch.nn.Module): + """With sub_group_size 1000 and ratio 0.5, `lin` is a sub-group kept on the GPU and `emb` a CPU sub-group.""" + + def __init__(self): + super().__init__() + self.lin = torch.nn.Linear(32, 32) + self.emb = torch.nn.Embedding(512, 32) + + def forward(self, tokens, inject): + loss = self.lin(self.emb(tokens)).float().pow(2).mean() + # inf * weight puts an inf gradient into `lin` only, so `emb` gets a finite gradient in the skipped step. + return loss + self.lin.weight[0, 0].float() * inject + + +class TestSuperOffloadSkippedStep(DistributedTest): + world_size = 1 + non_daemonic_procs = True + + def test_sub_group_with_finite_grads(self): + """A step skipped for an inf gradient elsewhere undoes the CPU sub-group's step; its Adam state must stay + valid. Each step looks up new embedding rows, which have zero Adam state before their first gradient.""" + masters = {} + for super_offload in (False, True): + config = _config(super_offload=super_offload, ratio=0.5, sub_group_size=1000, lr=1e-3, weight_decay=0.0) + torch.manual_seed(0) + model = InfInjectionModel() + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + masters[super_offload] = [] + for step in range(6): + tokens = (torch.arange(step * 64, step * 64 + 64).view(4, 16) % 512).to(engine.device) + loss = engine(tokens, float("inf") if step == 1 else 0.0) + engine.backward(loss) + engine.step() + masters[super_offload].append(_masters(engine)) + assert engine.skipped_steps == 1 + _close(engine) + for step, (a, e) in enumerate(zip(masters[True], masters[False])): + assert torch.isfinite(a).all(), f"step {step}" + # Undoing a step is exact up to rounding, which Adam amplifies for the rows whose state was zero. + torch.testing.assert_close(a, e, rtol=0, atol=1e-4, msg=f"step {step}") + + +class TestSuperOffloadTwoRanks(DistributedTest): + world_size = 2 + non_daemonic_procs = True + + def test_unused_param(self): + engine = _engine(_config(super_offload=False)) + expected = _train(engine, 6, _b_unused_on_even_steps("zero")) + _close(engine) + engine = _engine(_config(super_offload=True)) + actual = _train(engine, 6, _b_unused_on_even_steps("skip")) + _close(engine) + for step, (a, e) in enumerate(zip(actual, expected)): + assert torch.equal(a, e), f"step {step}: max diff {(a - e).abs().max().item()}" + + +def _resume_matches_uninterrupted(config, checkpoint_dir): + """Train 4 steps, save, train 4 more; a new engine loads the checkpoint and trains the same 4 steps.""" + engine = _engine(config) + _train(engine, 4, _all_used) + engine.save_checkpoint(checkpoint_dir, tag="t") + uninterrupted = _train(engine, 4, _all_used, first_step=4) + _close(engine) + engine = _engine(config) + engine.load_checkpoint(checkpoint_dir, tag="t") + resumed = _train(engine, 4, _all_used, first_step=4) + _close(engine) + for step, (r, u) in enumerate(zip(resumed, uninterrupted)): + assert torch.equal(r, u), f"step {step} after the resume: max diff {(r - u).abs().max().item()}" + + +class TestSuperOffloadCheckpoint(DistributedTest): + world_size = 2 + non_daemonic_procs = True + + # Betas of 0.5 make beta**t exact, so the resume must be bit-exact (CPU Adam recomputes beta**t on resume). + @pytest.mark.parametrize("super_offload,ratio", [(True, 1.0), (True, 0.5), (False, 0.5)]) + def test_resume(self, tmpdir, super_offload, ratio): + config = _config(super_offload=super_offload, ratio=ratio, sub_group_size=4500, betas=[0.5, 0.5]) + _resume_matches_uninterrupted(config, str(tmpdir)) + + +class TestPartialOffloadLRSchedule(DistributedTest): + world_size = 1 + non_daemonic_procs = True + + @pytest.mark.parametrize("super_offload", [True, False]) + def test_first_step_at_zero_lr(self, super_offload): + # With ratio < 1 the sub-groups kept on the device have their own optimizer; it must follow the scheduler. + config = _config(super_offload=super_offload, ratio=0.5, sub_group_size=4500) + config["scheduler"] = { + "type": "WarmupLR", + "params": { + "warmup_min_lr": 0.0, + "warmup_max_lr": 3e-3, + "warmup_num_steps": 4, + "warmup_type": "linear" + } + } + engine = _engine(config) + assert "cpu" in engine.optimizer.subgroup_to_device.values() + assert set(engine.optimizer.subgroup_to_device.values()) != {"cpu"} + before = _masters(engine) + after = _train(engine, 1, _all_used)[0] + _close(engine) + assert torch.equal(after, before) From 983cba3c21f44f85fa8974e696b27ebcf8a01592 Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Sun, 1 Nov 2026 05:25:47 +0000 Subject: [PATCH 02/10] SuperOffload: follow hyperparameter changes and loss scale increases, 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 --- csrc/adam/cpu_adam_impl.cpp | 4 +- deepspeed/runtime/engine.py | 4 + .../superoffload/superoffload_stage3.py | 58 +++++-- .../superoffload/superoffload_utils.py | 16 +- deepspeed/runtime/zero/stage3.py | 41 ----- tests/unit/v1/zero/test_invalid_grad_norm.py | 2 +- tests/unit/v1/zero/test_superoffload.py | 142 ++++++++++++------ 7 files changed, 157 insertions(+), 110 deletions(-) diff --git a/csrc/adam/cpu_adam_impl.cpp b/csrc/adam/cpu_adam_impl.cpp index 80adbd5f84d4..2adf981dad42 100644 --- a/csrc/adam/cpu_adam_impl.cpp +++ b/csrc/adam/cpu_adam_impl.cpp @@ -320,8 +320,8 @@ void adamw_rollback_inplace(float* params, momentum[i] = (momentum[i] - one_minus_beta1 * grad) / beta1; // exp_avg_sq is a sum of squares, so it is never negative. The inverse can leave a tiny // negative rounding residual (in about half of the elements that were 0 before the step), - // and sqrt() of it in the next step makes the parameter NaN. - variance[i] = std::max(0.0f, (variance[i] - one_minus_beta2 * grad * grad) / beta2); + // and sqrt() of it in the next step makes the parameter NaN. std::max(x, 0) keeps a NaN x. + variance[i] = std::max((variance[i] - one_minus_beta2 * grad * grad) / beta2, 0.0f); } --step_count; diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 93b72f3d012f..c73f59b3162a 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -4973,6 +4973,10 @@ def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, if save_latest and self.global_rank == 0: with open(os.path.join(save_dir, 'latest'), 'w') as fd: fd.write(tag) + elif self.super_offload(): + # The decoupled writer reads the fp32 weights and Adam state from shared memory, which SuperOffload + # updates in the next backward, before step() would commit this save. + self._commit_decoupled_checkpoint() dist.barrier() diff --git a/deepspeed/runtime/superoffload/superoffload_stage3.py b/deepspeed/runtime/superoffload/superoffload_stage3.py index ff767bc085bf..1fb2d1dd7f81 100644 --- a/deepspeed/runtime/superoffload/superoffload_stage3.py +++ b/deepspeed/runtime/superoffload/superoffload_stage3.py @@ -3,6 +3,7 @@ # DeepSpeed Team +import os import time from collections import defaultdict from typing import List @@ -57,7 +58,7 @@ def __init__( self.sub_group_to_param_num = {} # Per sub-group, the ds_ids of the params whose gradient for this step is in the fp32 gradient buffer. self.sub_group_grads_ready = defaultdict(set) - # Sub-groups stepped during backward, mapped to the (lr, grad_scale) of that step for a rollback. + # Sub-groups stepped during backward, mapped to the (hyperparameters, grad_scale) of that step for a rollback. self._submitted_cpu_sub_groups = {} self.async_cpuadam_num = 0 self.max_grad_numel = 0 @@ -87,6 +88,29 @@ def __init__( adamw_mode=adamw_mode, fp32_optimizer_states=fp32_optimizer_states, grad_dtype=self.master_weights_and_grads_dtype) + self._check_shm_size(fp32_optimizer_states) + + def _check_shm_size(self, fp32_optimizer_states): + """Warn when /dev/shm is too small for the tensors that SuperOffload shares with its worker process.""" + try: + stat = os.statvfs("/dev/shm") + except OSError: + return + numel = sum(self.fp32_partitioned_groups_flat[i].numel() for i, device in self.subgroup_to_device.items() + if device == 'cpu') + weight_size = torch.empty(0, dtype=self.master_weights_and_grads_dtype).element_size() + state_size = 4 if fp32_optimizer_states else weight_size + # The fp32 weights and gradients move there at the first step, the Adam state at the first checkpoint save + # or load. + needed = numel * 2 * (weight_size + state_size) + local_ranks = int(os.environ.get("LOCAL_WORLD_SIZE", os.environ.get("LOCAL_SIZE", 1))) + size = stat.f_blocks * stat.f_frsize + if size < needed * local_ranks: + logger.warning( + f"SuperOffload shares the fp32 weights, gradients and Adam state of the CPU sub-groups with " + f"its worker process through /dev/shm: {needed / 2**30:.1f} GiB on this rank, " + f"{local_ranks} rank(s) on this node, but /dev/shm has {size / 2**30:.1f} GiB. A checkpoint " + f"save or load may fail; increase the size of /dev/shm.") def _create_fp16_sub_groups(self, params_group): @@ -131,7 +155,6 @@ def step_with_gradscaler(optimizer): cur_device = self.subgroup_to_device[sub_group_id] if cur_device != 'cpu': - self._sync_backup_optimizer_group(param_group_id) self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] step_with_gradscaler(self.backup_optimizer) self.backup_optimizer.param_groups[param_group_id]['params'] = [] @@ -205,7 +228,7 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L if self.subgroup_to_device[i] == 'cpu' and self.speculative_cpu_step and self._grads_finite(i): param_group_id = self.sub_group_to_group_id[i] fp32_param = self.fp32_partitioned_groups_flat[i] - current_lr = self.optimizer.param_groups[param_group_id]['lr'] + hyperparams = self._hyperparams(param_group_id) # The gradients still carry the loss scale, which step() would remove. grad_scale = 1.0 / self.loss_scale @@ -213,9 +236,9 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L i, fp32_param.data, fp32_param.grad.data, - lr=current_lr, + hyperparams=hyperparams, grad_scale=grad_scale) - self._submitted_cpu_sub_groups[i] = (current_lr, grad_scale) + self._submitted_cpu_sub_groups[i] = (hyperparams, grad_scale) self.async_cpuadam_num += 1 result = self.superoffload_cpu_optimizer.get_result() @@ -261,8 +284,9 @@ def step(self, closure=None): timer_names.add(OPTIMIZER_STEP_TIMER) self.timers(OPTIMIZER_STEP_TIMER).start() - if self.check_clip_grads(scaled_global_grad_norm): - # The steps taken during backward used the unclipped gradients. + # The steps taken during backward used the unclipped gradients, and the loss scale of the backward, which + # _loss_scale_update_and_overflow_cleanup may have raised before unscale_and_clip_grads removes it. + if self.check_clip_grads(scaled_global_grad_norm) or self._loss_scale_changed(): self._rollback_submitted_cpu_sub_groups() for sub_group_id, group in enumerate(self.fp16_groups): @@ -274,12 +298,11 @@ def step(self, closure=None): elif sub_group_id not in self._submitted_cpu_sub_groups: param_group_id = self.sub_group_to_group_id[sub_group_id] fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - current_lr = self.optimizer.param_groups[param_group_id]['lr'] self._sync_cpu_optimizer_step(param_group_id, sub_group_id, fp32_param.data, fp32_param.grad.data, - lr=current_lr) + hyperparams=self._hyperparams(param_group_id)) self._reassign_or_swap_out_partitioned_parameters(sub_group_id) self._release_sub_group(sub_group_id, timer_names) @@ -327,6 +350,15 @@ def _grads_finite(self, sub_group_id): norms = [self.norm_for_param_grads[self.get_param_id(param)] for param in self.fp16_groups[sub_group_id]] return bool(torch.isfinite(torch.stack(norms)).all()) + def _hyperparams(self, param_group_id): + """The hyperparameters of a param group, which LR, momentum or weight decay schedules change.""" + group = self.optimizer.param_groups[param_group_id] + return {key: group[key] for key in ("lr", "betas", "eps", "weight_decay")} + + def _loss_scale_changed(self): + grad_scale = 1.0 / self.loss_scale + return any(scale != grad_scale for _, scale in self._submitted_cpu_sub_groups.values()) + def _wait_for_async_operations(self, timeout_seconds=60): """Wait for all pending asynchronous CPU optimizer operations to complete with timeout error. @@ -393,7 +425,7 @@ def _sync_cpu_optimizer_step(self, fp32_param_data, fp32_grad_data, rollback: bool = False, - lr: float = None, + hyperparams: dict = None, timeout_seconds: int = 60, grad_scale: float = 1.0): event_type = EventTypes.ROLLBACK if rollback else EventTypes.ADAM_STEP @@ -402,7 +434,7 @@ def _sync_cpu_optimizer_step(self, fp32_param_data, fp32_grad_data, rollback=rollback, - lr=lr, + hyperparams=hyperparams, grad_scale=grad_scale) # Wait for completion self._wait_for_single_async_result(event_type, timeout_seconds) @@ -413,7 +445,7 @@ def _rollback_submitted_cpu_sub_groups(self): self._rollback_cpu_sub_group(sub_group_id) def _rollback_cpu_sub_group(self, sub_group_id): - lr, grad_scale = self._submitted_cpu_sub_groups.pop(sub_group_id) + hyperparams, grad_scale = self._submitted_cpu_sub_groups.pop(sub_group_id) param_group_id = self.sub_group_to_group_id[sub_group_id] fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] self._sync_cpu_optimizer_step(param_group_id, @@ -421,7 +453,7 @@ def _rollback_cpu_sub_group(self, sub_group_id): fp32_param.data, fp32_param.grad.data, rollback=True, - lr=lr, + hyperparams=hyperparams, grad_scale=grad_scale) @instrument_w_nvtx diff --git a/deepspeed/runtime/superoffload/superoffload_utils.py b/deepspeed/runtime/superoffload/superoffload_utils.py index 83c924495063..a6e5dff072d6 100644 --- a/deepspeed/runtime/superoffload/superoffload_utils.py +++ b/deepspeed/runtime/superoffload/superoffload_utils.py @@ -24,7 +24,7 @@ class TaskKeys: PARAM_GROUP_ID = "param_group_id" SUB_GROUP_ID = "sub_group_id" ROLLBACK = "rollback" - LR = "lr" + HYPERPARAMS = "hyperparams" GRAD_SCALE = "grad_scale" COMMAND = "command" STATE = "state" @@ -144,7 +144,7 @@ def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, param_group_id = task[TaskKeys.PARAM_GROUP_ID] sub_group_id = task[TaskKeys.SUB_GROUP_ID] rollback = task.get(TaskKeys.ROLLBACK, False) - task_lr = task.get(TaskKeys.LR, None) + hyperparams = task.get(TaskKeys.HYPERPARAMS, None) grad_scale = task.get(TaskKeys.GRAD_SCALE, 1.0) logger.debug(f"Processing param_group_id: {param_group_id}, sub_group_id: {sub_group_id}") @@ -153,8 +153,8 @@ def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, del task[TaskKeys.PARAM_GRAD] task.clear() - if task_lr is not None: - optimizer.param_groups[param_group_id]['lr'] = task_lr + if hyperparams is not None: + optimizer.param_groups[param_group_id].update(hyperparams) grad_numel = param_grad.numel() if grad_numel > max_grad_numel: @@ -283,10 +283,12 @@ def async_step(self, fp32_param: torch.Tensor, fp32_grad: torch.Tensor, rollback: bool = False, - lr: float = None, + hyperparams: Dict[str, Any] = None, grad_scale: float = 1.0) -> None: """ Queue parameter for optimization in the worker process. + + hyperparams (lr, betas, eps, weight_decay) replace those of the param group in the worker before the step. """ if not self.cpuadam_process.is_alive(): raise RuntimeError("Worker process is not alive") @@ -298,8 +300,8 @@ def async_step(self, TaskKeys.SUB_GROUP_ID: sub_group_id, TaskKeys.ROLLBACK: rollback, } - if lr is not None: - task[TaskKeys.LR] = lr + if hyperparams is not None: + task[TaskKeys.HYPERPARAMS] = hyperparams if grad_scale != 1.0: task[TaskKeys.GRAD_SCALE] = grad_scale self.param_queue.put(task) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 4a1a5d699d85..27ddf8eec69b 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -1239,19 +1239,6 @@ def _create_fp16_sub_groups(self, params_group): return sub_groups - def _sync_backup_optimizer_group(self, param_group_id): - """Copy the hyperparameters of a param group to the optimizer of the sub-groups kept on the device. - - LR schedulers and set_lr() update self.optimizer only. The backup optimizer has its own first param group, - and after a checkpoint load its own copy of every group. - """ - backup_group = self.backup_optimizer.param_groups[param_group_id] - main_group = self.optimizer.param_groups[param_group_id] - if backup_group is main_group: - return - for key in ("lr", "betas", "eps", "weight_decay"): - backup_group[key] = main_group[key] - def _optimizer_step(self, sub_group_id): param_group_id = self.sub_group_to_group_id[sub_group_id] fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] @@ -1273,7 +1260,6 @@ def step_with_gradscaler(optimizer): step_with_gradscaler(self.optimizer) self.optimizer.param_groups[param_group_id]['params'] = [] else: - self._sync_backup_optimizer_group(param_group_id) self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] step_with_gradscaler(self.backup_optimizer) self.backup_optimizer.param_groups[param_group_id]['params'] = [] @@ -3454,20 +3440,6 @@ def _clear_fp32_optimizer_param_groups(self): for param_group in self.optimizer.param_groups: param_group['params'] = [] - def _set_backup_optimizer_param_groups(self): - # A new backup optimizer still holds its placeholder param, and after a load a group can share the param list - # of a replaced self.optimizer group. - self._clear_backup_optimizer_param_groups() - for sub_group_id, _ in enumerate(self.fp16_groups): - if self.subgroup_to_device[sub_group_id] != 'cpu': - param_group_id = self.sub_group_to_group_id[sub_group_id] - self.backup_optimizer.param_groups[param_group_id]['params'].append( - self.fp32_partitioned_groups_flat[sub_group_id]) - - def _clear_backup_optimizer_param_groups(self): - for param_group in self.backup_optimizer.param_groups: - param_group['params'] = [] - def _zero3_partition_group_metadata(self): metadata = [] has_autoep_expert_group = False @@ -3502,12 +3474,6 @@ def _rigid_state_dict(self): state_dict[FP32_FLAT_GROUPS] = self.fp32_partitioned_groups_flat self._clear_fp32_optimizer_param_groups() - # With partial offload, the sub-groups kept on the device have their Adam state in the backup optimizer. - if getattr(self, "backup_optimizer", None) is not None: - self._set_backup_optimizer_param_groups() - state_dict['backup_optimizer_state_dict'] = self.backup_optimizer.state_dict() - self._clear_backup_optimizer_param_groups() - return state_dict def state_dict(self): @@ -3617,13 +3583,6 @@ def _rigid_load_state_dict(self, state_dict, load_optimizer_states=True): self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) self._clear_fp32_optimizer_param_groups() - # Checkpoints saved without partial offload, or before its state was saved, do not have it. - has_backup_state = 'backup_optimizer_state_dict' in state_dict - if load_optimizer_states and has_backup_state and getattr(self, "backup_optimizer", None) is not None: - self._set_backup_optimizer_param_groups() - self.backup_optimizer.load_state_dict(state_dict['backup_optimizer_state_dict']) - self._clear_backup_optimizer_param_groups() - if self.swap_optimizer: # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint self.optimizer_swapper.purge_state() diff --git a/tests/unit/v1/zero/test_invalid_grad_norm.py b/tests/unit/v1/zero/test_invalid_grad_norm.py index 4d9bf0449166..152161fbc736 100644 --- a/tests/unit/v1/zero/test_invalid_grad_norm.py +++ b/tests/unit/v1/zero/test_invalid_grad_norm.py @@ -154,7 +154,7 @@ def test_bf16_optimizer_reports_invalid_step_and_recovers(monkeypatch): def test_superoffload_rolls_back_only_subgroups_submitted_this_step(): optimizer = object.__new__(SuperOffloadOptimizer_Stage3) - optimizer._submitted_cpu_sub_groups = {1: (0.1, 1.0)} + optimizer._submitted_cpu_sub_groups = {1: ({"lr": 0.1}, 1.0)} optimizer.sub_group_to_group_id = {0: 10, 1: 11} parameter = SimpleNamespace(data=torch.tensor([1.0]), grad=SimpleNamespace(data=torch.tensor([2.0]))) optimizer.fp32_partitioned_groups_flat = [parameter, parameter] diff --git a/tests/unit/v1/zero/test_superoffload.py b/tests/unit/v1/zero/test_superoffload.py index 06a91093d0ad..9408dace59b9 100644 --- a/tests/unit/v1/zero/test_superoffload.py +++ b/tests/unit/v1/zero/test_superoffload.py @@ -65,14 +65,26 @@ def _config(super_offload, clip=0.0, ga=1, precision="bf16", ratio=1.0, sub_grou if precision == "fp16": config["fp16"] = {"enabled": True, "initial_scale_power": 24} elif precision != "fp32": - config["bf16"] = {"enabled": True, "bf16_master_weights_and_grads": precision == "bf16_master"} + config["bf16"] = { + "enabled": True, + "bf16_master_weights_and_grads": precision in ("bf16_master", "bf16_states"), + "bf16_optimizer_states": precision == "bf16_states", + } return config -def _engine(config): +def _engine(config, two_groups=False): torch.manual_seed(0) model = PartlyUsedModel() - engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + params = model.parameters() + if two_groups: + # Weights with weight decay and biases without, as most training scripts set them up. + named = list(model.named_parameters()) + params = [ + dict(params=[p for n, p in named if n.endswith("weight")]), + dict(params=[p for n, p in named if n.endswith("bias")], weight_decay=0.0) + ] + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=params, config=config) return engine @@ -85,10 +97,11 @@ def _close(engine): engine.optimizer.superoffload_cpu_optimizer.close() -def _train(engine, steps, uses_for, first_step=0): +def _train(engine, steps, uses_for, first_step=0, before_step=None): """Train `steps` optimizer steps and return the master weights after each one. - uses_for(step, micro_step) gives the `uses` of PartlyUsedModel.forward. + uses_for(step, micro_step) gives the `uses` of PartlyUsedModel.forward. before_step(engine, step), if given, runs + before each step, as an LR or momentum schedule would. """ dtype = torch.float32 if engine.fp16_enabled(): @@ -98,6 +111,8 @@ def _train(engine, steps, uses_for, first_step=0): generator = torch.Generator().manual_seed(1 + first_step + 100 * dist.get_rank()) masters = [] for step in range(first_step, first_step + steps): + if before_step is not None: + before_step(engine, step) for micro_step in range(engine.gradient_accumulation_steps()): x = torch.randn(4, HIDDEN, generator=generator).to(engine.device, dtype) loss = engine(x, uses_for(step, micro_step)) @@ -129,6 +144,8 @@ def _b_and_c_in_one_micro_step(unused): "adam": (dict(type="Adam", adam_w_mode=False), lambda unused: _all_used, True), "no_bias_correction": (dict(bias_correction=False), lambda unused: _all_used, True), "bf16_master_weights": (dict(precision="bf16_master"), lambda unused: _all_used, True), + "bf16_optimizer_states": (dict(precision="bf16_states"), lambda unused: _all_used, True), + "two_param_groups": (dict(two_groups=True, sub_group_size=4500), _b_unused_on_even_steps, True), # Clipping undoes the steps taken during backward and redoes them, which is exact up to rounding. In fp32 the # rounding stays small; bf16 weights would amplify it. "clipping": (dict(clip=0.05, precision="fp32"), lambda unused: _all_used, False), @@ -145,21 +162,80 @@ class TestSuperOffloadMatchesZeRO3(DistributedTest): @pytest.mark.parametrize("case", list(CASES)) def test(self, case): overrides, uses_for, exact = CASES[case] + overrides = dict(overrides) + two_groups = overrides.pop("two_groups", False) # The reference gives an unused param an exactly zero gradient, which SuperOffload must give a param that # gets no gradient (as ZeRO-1/2 do). - engine = _engine(_config(super_offload=False, **overrides)) + engine = _engine(_config(super_offload=False, **overrides), two_groups) expected = _train(engine, 8, uses_for("zero")) _close(engine) - engine = _engine(_config(super_offload=True, **overrides)) + engine = _engine(_config(super_offload=True, **overrides), two_groups) actual = _train(engine, 8, uses_for("skip")) _close(engine) + _assert_same(actual, expected, exact) - for step, (a, e) in enumerate(zip(actual, expected)): - assert torch.isfinite(a).all(), f"step {step}" - if exact: - assert torch.equal(a, e), f"step {step}: max diff {(a - e).abs().max().item()}" - else: - torch.testing.assert_close(a, e, rtol=0, atol=1e-5, msg=f"step {step}") + +def _assert_same(actual, expected, exact, atol=1e-5): + for step, (a, e) in enumerate(zip(actual, expected)): + assert torch.isfinite(a).all(), f"step {step}" + if exact: + assert torch.equal(a, e), f"step {step}: max diff {(a - e).abs().max().item()}" + else: + torch.testing.assert_close(a, e, rtol=0, atol=atol, msg=f"step {step}") + + +def _change_hyperparams(engine, step): + for group in engine.optimizer.param_groups: + group["betas"] = (0.9 - 0.05 * step, 0.999 - 0.01 * step) + group["eps"] = 1e-8 * 10**step + group["weight_decay"] = 0.01 + 0.1 * step + + +class TestSuperOffloadFollowsZeRO3(DistributedTest): + world_size = 1 + non_daemonic_procs = True + + def test_hyperparameter_changes(self): + """Momentum and weight decay schedules change betas and weight_decay of the param groups between steps.""" + results = [] + for super_offload in (False, True): + engine = _engine(_config(super_offload=super_offload), two_groups=True) + results.append(_train(engine, 6, _all_used, before_step=_change_hyperparams)) + _close(engine) + _assert_same(results[1], results[0], exact=True) + + def test_fp16_loss_scale_increase(self): + """Dynamic loss scaling raises the scale in step(), after the CPU steps taken during backward used the old + one; those steps are undone and redone, exact up to rounding, which fp16 weights amplify (1.4e-5 after 6 + steps; 3.9e-3 when the steps are kept).""" + results = [] + for super_offload in (False, True): + config = _config(super_offload=super_offload, precision="fp16") + config["fp16"] = {"enabled": True, "initial_scale_power": 8, "loss_scale_window": 2, "hysteresis": 1} + engine = _engine(config) + results.append(_train(engine, 6, _all_used)) + assert engine.optimizer.loss_scale > 2**8 + _close(engine) + _assert_same(results[1], results[0], exact=False, atol=1e-4) + + def test_cpu_steps_run_during_backward(self): + engine = _engine(_config(super_offload=True)) + cpu_optimizer = engine.optimizer.superoffload_cpu_optimizer + submitted = [] + async_step = cpu_optimizer.async_step + + def recording_async_step(*args, **kwargs): + submitted.append(kwargs.get("rollback", False)) + async_step(*args, **kwargs) + + cpu_optimizer.async_step = recording_async_step + x = torch.randn(4, HIDDEN).to(engine.device, torch.bfloat16) + engine.backward(engine(x, {})) + cpu_sub_groups = list(engine.optimizer.subgroup_to_device.values()).count("cpu") + assert submitted == [False] * cpu_sub_groups + engine.step() + assert submitted == [False] * cpu_sub_groups + _close(engine) class InfInjectionModel(torch.nn.Module): @@ -219,14 +295,14 @@ def test_unused_param(self): assert torch.equal(a, e), f"step {step}: max diff {(a - e).abs().max().item()}" -def _resume_matches_uninterrupted(config, checkpoint_dir): +def _resume_matches_uninterrupted(config, checkpoint_dir, two_groups): """Train 4 steps, save, train 4 more; a new engine loads the checkpoint and trains the same 4 steps.""" - engine = _engine(config) + engine = _engine(config, two_groups) _train(engine, 4, _all_used) engine.save_checkpoint(checkpoint_dir, tag="t") uninterrupted = _train(engine, 4, _all_used, first_step=4) _close(engine) - engine = _engine(config) + engine = _engine(config, two_groups) engine.load_checkpoint(checkpoint_dir, tag="t") resumed = _train(engine, 4, _all_used, first_step=4) _close(engine) @@ -239,33 +315,7 @@ class TestSuperOffloadCheckpoint(DistributedTest): non_daemonic_procs = True # Betas of 0.5 make beta**t exact, so the resume must be bit-exact (CPU Adam recomputes beta**t on resume). - @pytest.mark.parametrize("super_offload,ratio", [(True, 1.0), (True, 0.5), (False, 0.5)]) - def test_resume(self, tmpdir, super_offload, ratio): - config = _config(super_offload=super_offload, ratio=ratio, sub_group_size=4500, betas=[0.5, 0.5]) - _resume_matches_uninterrupted(config, str(tmpdir)) - - -class TestPartialOffloadLRSchedule(DistributedTest): - world_size = 1 - non_daemonic_procs = True - - @pytest.mark.parametrize("super_offload", [True, False]) - def test_first_step_at_zero_lr(self, super_offload): - # With ratio < 1 the sub-groups kept on the device have their own optimizer; it must follow the scheduler. - config = _config(super_offload=super_offload, ratio=0.5, sub_group_size=4500) - config["scheduler"] = { - "type": "WarmupLR", - "params": { - "warmup_min_lr": 0.0, - "warmup_max_lr": 3e-3, - "warmup_num_steps": 4, - "warmup_type": "linear" - } - } - engine = _engine(config) - assert "cpu" in engine.optimizer.subgroup_to_device.values() - assert set(engine.optimizer.subgroup_to_device.values()) != {"cpu"} - before = _masters(engine) - after = _train(engine, 1, _all_used)[0] - _close(engine) - assert torch.equal(after, before) + @pytest.mark.parametrize("two_groups", [False, True]) + def test_resume(self, tmpdir, two_groups): + config = _config(super_offload=True, sub_group_size=4500, betas=[0.5, 0.5]) + _resume_matches_uninterrupted(config, str(tmpdir), two_groups) From afded1772ac0a1b889f0eeac97aa80bcce90b653 Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Sun, 1 Nov 2026 05:25:55 +0000 Subject: [PATCH 03/10] ZeRO-3 partial offload: follow hyperparameter schedules on the device 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 --- .../superoffload/superoffload_stage3.py | 1 + deepspeed/runtime/zero/stage3.py | 94 +++++++- .../unit/v1/zero/test_zero_partial_offload.py | 217 ++++++++++++++++++ 3 files changed, 300 insertions(+), 12 deletions(-) create mode 100644 tests/unit/v1/zero/test_zero_partial_offload.py diff --git a/deepspeed/runtime/superoffload/superoffload_stage3.py b/deepspeed/runtime/superoffload/superoffload_stage3.py index d6a4270f5aaa..aee0b4d86b27 100644 --- a/deepspeed/runtime/superoffload/superoffload_stage3.py +++ b/deepspeed/runtime/superoffload/superoffload_stage3.py @@ -103,6 +103,7 @@ def step_with_gradscaler(optimizer): cur_device = self.subgroup_to_device[sub_group_id] if cur_device != 'cpu': + self._sync_backup_optimizer_group(param_group_id) self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] step_with_gradscaler(self.backup_optimizer) self.backup_optimizer.param_groups[param_group_id]['params'] = [] diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 27ddf8eec69b..0fc616384316 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -314,16 +314,25 @@ def _enforce_optimizer_offload(): backup_gpu_tensor = torch.randn(1, device=get_accelerator().device_name()).to(self.dtype) backup_gpu_param = torch.nn.Parameter(backup_gpu_tensor) assert type(init_optimizer) == DeepSpeedCPUAdam, 'Hybrid Optimizer Only Supports DeepSpeedCPUAdam' - self.backup_optimizer = torch.optim.AdamW([backup_gpu_param], - lr=self.optimizer.param_groups[0]["lr"], - betas=self.optimizer.param_groups[0]["betas"], - eps=self.optimizer.param_groups[0]["eps"], - weight_decay=self.optimizer.param_groups[0]["weight_decay"], - amsgrad=self.optimizer.param_groups[0]["amsgrad"]) - # Multiple param_groups configs for back-up optimizer - if len(self.optimizer.param_groups) > 1: - for i in range(1, len(self.optimizer.param_groups)): - self.backup_optimizer.add_param_group(self.optimizer.param_groups[i]) + if not all(group.get("bias_correction", True) for group in self.optimizer.param_groups): + raise ValueError( + "offload_optimizer.ratio < 1 steps the sub-groups kept on the device with torch Adam, " + "which always applies bias correction; bias_correction=False is not supported.") + # The same update as DeepSpeedCPUAdam: decoupled weight decay, or L2 when adam_w_mode is off. + backup_optimizer_class = torch.optim.AdamW if init_optimizer.adam_w_mode else torch.optim.Adam + self.backup_optimizer = backup_optimizer_class([backup_gpu_param], + lr=self.optimizer.param_groups[0]["lr"], + betas=self.optimizer.param_groups[0]["betas"], + eps=self.optimizer.param_groups[0]["eps"], + weight_decay=self.optimizer.param_groups[0]["weight_decay"], + amsgrad=self.optimizer.param_groups[0]["amsgrad"]) + # Each param group gets its own dict; _sync_backup_optimizer_group copies the hyperparameters before a + # step. A dict shared with self.optimizer would also get the CPU sub-groups that + # _set_fp32_optimizer_param_groups adds for a checkpoint, and keep them after a load replaces + # self.optimizer's dicts. + for group in self.optimizer.param_groups[1:]: + backup_group = {key: value for key, value in group.items() if key != "params"} + self.backup_optimizer.add_param_group(dict(backup_group, params=[])) self._initialize_zenflow_stage3_epilogue(zenflow_config, overlap_comm) @@ -1239,6 +1248,23 @@ def _create_fp16_sub_groups(self, params_group): return sub_groups + def _sync_backup_optimizer_group(self, param_group_id): + """Copy the hyperparameters of a param group to the optimizer of the sub-groups kept on the device. + + LR, momentum and weight decay schedules, set_lr() and checkpoint loads change self.optimizer's groups only. + """ + backup_group = self.backup_optimizer.param_groups[param_group_id] + main_group = self.optimizer.param_groups[param_group_id] + for key in ("lr", "betas", "eps", "weight_decay"): + backup_group[key] = main_group[key] + + def _sub_group_optimizer(self, sub_group_id): + """The optimizer that steps a sub-group and holds its state.""" + if (self.offload_optimizer and self.subgroup_to_device[sub_group_id] != 'cpu' + and not (self.use_muon and self.sub_groups_using_muon[sub_group_id])): + return self.backup_optimizer + return self.optimizer + def _optimizer_step(self, sub_group_id): param_group_id = self.sub_group_to_group_id[sub_group_id] fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] @@ -1260,6 +1286,7 @@ def step_with_gradscaler(optimizer): step_with_gradscaler(self.optimizer) self.optimizer.param_groups[param_group_id]['params'] = [] else: + self._sync_backup_optimizer_group(param_group_id) self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] step_with_gradscaler(self.backup_optimizer) self.backup_optimizer.param_groups[param_group_id]['params'] = [] @@ -3190,7 +3217,8 @@ def _get_fp32_opt_state_partition(self, param, release_swap_buffers, optim_state if optim_state_key is None: fp32_opt_state = fp32_param.narrow(0, dest_offset, num_elements) else: - fp32_opt_state = self.optimizer.state[fp32_param][optim_state_key].narrow(0, dest_offset, num_elements) + optimizer = self._sub_group_optimizer(group_idx) + fp32_opt_state = optimizer.state[fp32_param][optim_state_key].narrow(0, dest_offset, num_elements) if self._swappable_optimizer_subgroup(group_idx) and release_swap_buffers: self._release_swap_buffers(group_idx) @@ -3470,12 +3498,52 @@ def _rigid_state_dict(self): state_dict['ds_zero_partition_groups'] = zero3_partition_group_metadata self._set_fp32_optimizer_param_groups() - state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + # The sub-groups kept on the device have their state in the backup optimizer. Save it with the other + # sub-groups', as DeepSpeedCPUAdam keeps it, so that the checkpoint does not depend on offload_optimizer.ratio. + backup_state_params = self._add_backup_optimizer_state() + try: + state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + finally: + for fp32_param in backup_state_params: + del self.optimizer.state[fp32_param] state_dict[FP32_FLAT_GROUPS] = self.fp32_partitioned_groups_flat self._clear_fp32_optimizer_param_groups() return state_dict + def _add_backup_optimizer_state(self): + """Put the Adam state of the sub-groups kept on the device into self.optimizer.state; return their params.""" + added = [] + if getattr(self, "backup_optimizer", None) is None: + return added + for sub_group_id, fp32_param in enumerate(self.fp32_partitioned_groups_flat): + state = self.backup_optimizer.state.get(fp32_param) + if state and self._sub_group_optimizer(sub_group_id) is self.backup_optimizer: + self.optimizer.state[fp32_param] = { + "step": int(state["step"]), + "exp_avg": state["exp_avg"], + "exp_avg_sq": state["exp_avg_sq"], + } + added.append(fp32_param) + return added + + def _move_state_to_backup_optimizer(self): + """Move the loaded Adam state of the sub-groups kept on the device to the backup optimizer.""" + if getattr(self, "backup_optimizer", None) is None: + return + for sub_group_id, fp32_param in enumerate(self.fp32_partitioned_groups_flat): + if self._sub_group_optimizer(sub_group_id) is not self.backup_optimizer: + continue + state = self.optimizer.state.pop(fp32_param, None) + self.backup_optimizer.state.pop(fp32_param, None) + if state: + # torch Adam keeps the step count as a float32 tensor, DeepSpeedCPUAdam as an int. + self.backup_optimizer.state[fp32_param] = { + "step": torch.tensor(float(state["step"])), + "exp_avg": state["exp_avg"], + "exp_avg_sq": state["exp_avg_sq"], + } + def state_dict(self): """ Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. @@ -3582,6 +3650,7 @@ def _rigid_load_state_dict(self, state_dict, load_optimizer_states=True): self._set_fp32_optimizer_param_groups() self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) self._clear_fp32_optimizer_param_groups() + self._move_state_to_backup_optimizer() if self.swap_optimizer: # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint @@ -3712,6 +3781,7 @@ def load_hp_checkpoint_state_from_checkpoint_dir_stage3(self, checkpoint_dir): self.optimizer.state[fp32_param]['step'] = state_step else: self.optimizer.state[fp32_param][key] = key_tensor + self._move_state_to_backup_optimizer() for param_group in self.optimizer.param_groups: # Generally, the hyperparameters of each parameter should be the same, we can obtain from any parameter. diff --git a/tests/unit/v1/zero/test_zero_partial_offload.py b/tests/unit/v1/zero/test_zero_partial_offload.py new file mode 100644 index 000000000000..162a662fb4f2 --- /dev/null +++ b/tests/unit/v1/zero/test_zero_partial_offload.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""ZeRO-3 with offload_optimizer.ratio < 1, where the sub-groups kept on the device are stepped by a torch Adam.""" + +import glob +import os + +import pytest +import torch + +import deepspeed +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import CPUAdamBuilder +from deepspeed.utils import safe_get_full_optimizer_state +from unit.common import DistributedTest + +if not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: + pytest.skip("cpu-adam is not compatible", allow_module_level=True) + +HIDDEN = 64 + + +class Net(torch.nn.Module): + + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList(torch.nn.Linear(HIDDEN, HIDDEN) for _ in range(3)) + + def forward(self, x): + for layer in self.layers: + x = torch.nn.functional.gelu(layer(x)) + return x.float().pow(2).mean() + + +def _config(ratio, super_offload=False, **optimizer_params): + optimizer_params.setdefault("lr", 3e-3) + optimizer_params.setdefault("weight_decay", 0.01) + # 3 sub-groups: the first two weights, the last weight, the biases. Ratio 0.5 keeps the first on the device. + return { + "train_micro_batch_size_per_gpu": 4, + "steps_per_print": 1000, + "optimizer": { + "type": "AdamW", + "params": optimizer_params + }, + "zero_optimization": { + "stage": 3, + "sub_group_size": 4500, + "offload_optimizer": { + "device": "cpu", + "ratio": ratio, + "super_offload": super_offload + }, + }, + "bf16": { + "enabled": True + }, + } + + +def _engine(config): + torch.manual_seed(0) + model = Net() + named = list(model.named_parameters()) + # Weights with weight decay and biases without, as most training scripts set them up. + params = [ + dict(params=[p for n, p in named if n.endswith("weight")]), + dict(params=[p for n, p in named if n.endswith("bias")], weight_decay=0.0) + ] + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=params, config=config) + return engine + + +def _train(engine, steps, first_step=0): + generator = torch.Generator().manual_seed(1 + first_step + 100 * dist.get_rank()) + dtype = torch.bfloat16 if engine.bfloat16_enabled() else torch.float32 + masters = [] + for _ in range(steps): + x = torch.randn(4, HIDDEN, generator=generator).to(engine.device, dtype) + loss = engine(x) + engine.backward(loss) + engine.step() + masters.append( + torch.cat([t.detach().float().flatten().cpu() for t in engine.optimizer.fp32_partitioned_groups_flat])) + return masters + + +def _close(engine): + if hasattr(engine.optimizer, "superoffload_cpu_optimizer"): + engine.optimizer.superoffload_cpu_optimizer.close() + + +def _device_params(engine): + optimizer = engine.optimizer + return [optimizer.fp32_partitioned_groups_flat[i] for i, d in optimizer.subgroup_to_device.items() if d != "cpu"] + + +def _adam_state(engine): + return { + (name, key): safe_get_full_optimizer_state(p, key).cpu() + for name, p in engine.module.named_parameters() + for key in ("exp_avg", "exp_avg_sq") + } + + +class TestPartialOffloadHyperparameters(DistributedTest): + world_size = 1 + # SuperOffload starts its CPU worker process, which a daemonic test process cannot do. + non_daemonic_procs = True + + @pytest.mark.parametrize("super_offload", [False, True]) + def test_first_step_at_zero_lr(self, super_offload): + if super_offload and get_accelerator().device_name() != "cuda": + pytest.skip("SuperOffload supports only CUDA") + config = _config(ratio=0.5, super_offload=super_offload) + config["scheduler"] = { + "type": "WarmupLR", + "params": { + "warmup_min_lr": 0.0, + "warmup_max_lr": 3e-3, + "warmup_num_steps": 4, + "warmup_type": "linear" + } + } + engine = _engine(config) + assert _device_params(engine) and "cpu" in engine.optimizer.subgroup_to_device.values() + before = torch.cat([t.detach().float().flatten().cpu() for t in engine.optimizer.fp32_partitioned_groups_flat]) + after = _train(engine, 1)[0] + assert torch.equal(after, before) + # A momentum or weight decay schedule changes the other hyperparameters of the param groups. + for group in engine.optimizer.param_groups: + group["betas"], group["eps"], group["weight_decay"] = (0.8, 0.99), 1e-6, 0.2 + _train(engine, 1) + optimizer = engine.optimizer + # (The scheduler has already set the lr of the next step.) + for i in {optimizer.sub_group_to_group_id[i] for i, d in optimizer.subgroup_to_device.items() if d != "cpu"}: + for key in ("betas", "eps", "weight_decay"): + assert optimizer.backup_optimizer.param_groups[i][key] == optimizer.param_groups[i][key], key + _close(engine) + + +class TestPartialOffloadCheckpoint(DistributedTest): + world_size = 2 + + def test_resume(self, tmpdir): + """Betas of 0.5 make beta**t exact, so the resume is bit-exact (CPU Adam recomputes beta**t on resume).""" + config = _config(ratio=0.5, betas=[0.5, 0.5]) + engine = _engine(config) + _train(engine, 4) + engine.save_checkpoint(str(tmpdir), tag="t") + uninterrupted = _train(engine, 4, first_step=4) + engine = _engine(config) + engine.load_checkpoint(str(tmpdir), tag="t") + resumed = _train(engine, 4, first_step=4) + for step, (r, u) in enumerate(zip(resumed, uninterrupted)): + assert torch.equal(r, u), f"step {step} after the resume: max diff {(r - u).abs().max().item()}" + + +class TestPartialOffloadCheckpointPortability(DistributedTest): + world_size = 1 + + @pytest.mark.parametrize("load_ratio", [0.25, 0.75, 1.0]) + def test_load_with_another_ratio(self, tmpdir, load_ratio): + engine = _engine(_config(ratio=0.5)) + _train(engine, 3) + saved = _adam_state(engine) + engine.save_checkpoint(str(tmpdir), tag="t") + engine = _engine(_config(ratio=load_ratio)) + engine.load_checkpoint(str(tmpdir), tag="t") + loaded = _adam_state(engine) + for key, value in saved.items(): + assert torch.equal(loaded[key], value), key + _train(engine, 1) + engine.save_checkpoint(str(tmpdir), tag="again") + + def test_load_checkpoint_without_device_state(self, tmpdir): + """Checkpoints saved before the state of the sub-groups kept on the device was saved have none for them.""" + engine = _engine(_config(ratio=0.5)) + _train(engine, 3) + engine.save_checkpoint(str(tmpdir), tag="t") + device_ids = [i for i, d in engine.optimizer.subgroup_to_device.items() if d != "cpu"] + path, = glob.glob(os.path.join(str(tmpdir), "t", "*zero_pp_rank_0_mp_rank_00_optim_states.pt")) + checkpoint = torch.load(path, weights_only=False) + for i in device_ids: + checkpoint["optimizer_state_dict"]["optimizer_state_dict"]["state"].pop(i, None) + torch.save(checkpoint, path) + + engine = _engine(_config(ratio=0.5)) + engine.load_checkpoint(str(tmpdir), tag="t") + _train(engine, 1) + # Only the sub-groups kept on the device are stepped by the backup optimizer. + assert set(engine.optimizer.backup_optimizer.state) == set(_device_params(engine)) + engine.save_checkpoint(str(tmpdir), tag="again") + + +class TestPartialOffloadAdamSettings(DistributedTest): + world_size = 1 + + def test_l2_weight_decay(self): + """With adam_w_mode off, the sub-groups kept on the device apply L2 weight decay, as DeepSpeedCPUAdam does.""" + results = [] + for ratio in (1.0, 0.5): + config = _config(ratio=ratio, weight_decay=0.1) + config["optimizer"] = {"type": "Adam", "params": {**config["optimizer"]["params"], "adam_w_mode": False}} + del config["bf16"] + engine = _engine(config) + results.append(_train(engine, 4)) + for step, (partial, cpu) in enumerate(zip(results[1], results[0])): + # torch Adam and DeepSpeedCPUAdam round differently. + torch.testing.assert_close(partial, cpu, rtol=0, atol=1e-6, msg=f"step {step}") + + def test_bias_correction_off_is_rejected(self): + config = _config(ratio=0.5) + config["optimizer"]["params"]["bias_correction"] = False + with pytest.raises(ValueError, match="bias_correction"): + _engine(config) From d53ffb6a9aeff4b9e1da89b6c8a7004d81c5e5af Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Sun, 1 Nov 2026 06:39:39 +0000 Subject: [PATCH 04/10] ZeRO-3 partial offload: load the device sub-groups' moments in the param 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 --- deepspeed/runtime/zero/stage3.py | 7 ++++--- tests/unit/v1/zero/test_zero_partial_offload.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 0fc616384316..e537b6ba4910 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -3537,11 +3537,12 @@ def _move_state_to_backup_optimizer(self): state = self.optimizer.state.pop(fp32_param, None) self.backup_optimizer.state.pop(fp32_param, None) if state: - # torch Adam keeps the step count as a float32 tensor, DeepSpeedCPUAdam as an int. + # torch Adam keeps the moments in the param dtype and the step count as a float32 tensor; + # DeepSpeedCPUAdam keeps an int step and, with fp32_optimizer_states, fp32 moments for bf16 params. self.backup_optimizer.state[fp32_param] = { "step": torch.tensor(float(state["step"])), - "exp_avg": state["exp_avg"], - "exp_avg_sq": state["exp_avg_sq"], + "exp_avg": state["exp_avg"].to(fp32_param.dtype), + "exp_avg_sq": state["exp_avg_sq"].to(fp32_param.dtype), } def state_dict(self): diff --git a/tests/unit/v1/zero/test_zero_partial_offload.py b/tests/unit/v1/zero/test_zero_partial_offload.py index 162a662fb4f2..9b2c98cf1d79 100644 --- a/tests/unit/v1/zero/test_zero_partial_offload.py +++ b/tests/unit/v1/zero/test_zero_partial_offload.py @@ -174,6 +174,21 @@ def test_load_with_another_ratio(self, tmpdir, load_ratio): _train(engine, 1) engine.save_checkpoint(str(tmpdir), tag="again") + def test_resume_with_bf16_master_weights(self, tmpdir): + """torch Adam keeps the moments of the device sub-groups in the param dtype, also after a load that returns + fp32 moments for bf16 params (DeepSpeedCPUAdam with fp32_optimizer_states).""" + config = _config(ratio=0.5) + config["bf16"]["bf16_master_weights_and_grads"] = True + engine = _engine(config) + _train(engine, 2) + engine.save_checkpoint(str(tmpdir), tag="t") + engine = _engine(config) + engine.load_checkpoint(str(tmpdir), tag="t") + assert torch.isfinite(_train(engine, 1)[0]).all() + for fp32_param in _device_params(engine): + state = engine.optimizer.backup_optimizer.state[fp32_param] + assert state["exp_avg"].dtype == state["exp_avg_sq"].dtype == fp32_param.dtype + def test_load_checkpoint_without_device_state(self, tmpdir): """Checkpoints saved before the state of the sub-groups kept on the device was saved have none for them.""" engine = _engine(_config(ratio=0.5)) From 9867efa219c2acfcd460204fde5268b0f451552d Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Sun, 1 Nov 2026 07:47:13 +0000 Subject: [PATCH 05/10] SuperOffload: do not step during backward when a beta is 0 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 --- deepspeed/runtime/superoffload/superoffload_stage3.py | 10 ++++++++-- tests/unit/v1/zero/test_superoffload.py | 4 ++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/superoffload/superoffload_stage3.py b/deepspeed/runtime/superoffload/superoffload_stage3.py index 1fb2d1dd7f81..7d1f1c0dfead 100644 --- a/deepspeed/runtime/superoffload/superoffload_stage3.py +++ b/deepspeed/runtime/superoffload/superoffload_stage3.py @@ -225,8 +225,9 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L get_accelerator().current_stream().synchronize() # A param that gets a second gradient adds its sub-group again. for i in dict.fromkeys(completed_sub_groups): - if self.subgroup_to_device[i] == 'cpu' and self.speculative_cpu_step and self._grads_finite(i): - param_group_id = self.sub_group_to_group_id[i] + param_group_id = self.sub_group_to_group_id[i] + if (self.subgroup_to_device[i] == 'cpu' and self.speculative_cpu_step + and self._betas_nonzero(param_group_id) and self._grads_finite(i)): fp32_param = self.fp32_partitioned_groups_flat[i] hyperparams = self._hyperparams(param_group_id) # The gradients still carry the loss scale, which step() would remove. @@ -350,6 +351,11 @@ def _grads_finite(self, sub_group_id): norms = [self.norm_for_param_grads[self.get_param_id(param)] for param in self.fp16_groups[sub_group_id]] return bool(torch.isfinite(torch.stack(norms)).all()) + def _betas_nonzero(self, param_group_id): + """Whether a step of the param group can be rolled back: the rollback divides the moments by beta1 and beta2. + Checked at each step, as momentum schedules change the betas.""" + return all(beta != 0 for beta in self.optimizer.param_groups[param_group_id]["betas"]) + def _hyperparams(self, param_group_id): """The hyperparameters of a param group, which LR, momentum or weight decay schedules change.""" group = self.optimizer.param_groups[param_group_id] diff --git a/tests/unit/v1/zero/test_superoffload.py b/tests/unit/v1/zero/test_superoffload.py index 9408dace59b9..bbcb80970254 100644 --- a/tests/unit/v1/zero/test_superoffload.py +++ b/tests/unit/v1/zero/test_superoffload.py @@ -149,6 +149,10 @@ def _b_and_c_in_one_micro_step(unused): # Clipping undoes the steps taken during backward and redoes them, which is exact up to rounding. In fp32 the # rounding stays small; bf16 weights would amplify it. "clipping": (dict(clip=0.05, precision="fp32"), lambda unused: _all_used, False), + # A step cannot be undone with a beta of 0 (the rollback divides by the betas), so those CPU steps wait for + # step() and stay exact. + "clipping_beta1_zero": (dict(clip=0.05, precision="fp32", betas=[0.0, 0.999]), lambda unused: _all_used, True), + "clipping_beta2_zero": (dict(clip=0.05, precision="fp32", betas=[0.9, 0.0]), lambda unused: _all_used, True), # Overflow steps are skipped; the steps taken during backward are undone, exact up to rounding. "fp16_overflow": (dict(precision="fp16", sub_group_size=4500), lambda unused: _all_used, False), } From 68d10865e2a571aed81f6c926cf84830bea3f0b4 Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Thu, 5 Nov 2026 22:23:48 +0000 Subject: [PATCH 06/10] CPU Adam: release the GIL during steps, step fp32 params on bf16 gradients 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 --- csrc/adam/cpu_adam.cpp | 12 ++++- csrc/adam/cpu_adam_impl.cpp | 50 +++++++++++++------- csrc/includes/cpu_adam.h | 67 ++++++++++++++++++--------- tests/unit/ops/adam/test_cpu_adam.py | 68 ++++++++++++++++++++++++++++ 4 files changed, 156 insertions(+), 41 deletions(-) diff --git a/csrc/adam/cpu_adam.cpp b/csrc/adam/cpu_adam.cpp index 72b7b2a97473..90c7c188075b 100644 --- a/csrc/adam/cpu_adam.cpp +++ b/csrc/adam/cpu_adam.cpp @@ -8,8 +8,16 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { using namespace pybind11::literals; - m.def("adam_update", &ds_adam_step, "DeepSpeed CPU Adam update (C++)"); - m.def("adam_rollback", &ds_adam_rollback, "DeepSpeed CPU Adam rollback (C++)"); + // The update and the rollback release the GIL, so that a CPU Adam step in another thread lets + // the training thread run. + m.def("adam_update", + &ds_adam_step, + "DeepSpeed CPU Adam update (C++)", + pybind11::call_guard()); + m.def("adam_rollback", + &ds_adam_rollback, + "DeepSpeed CPU Adam rollback (C++)", + pybind11::call_guard()); m.def("create_adam", &create_adam_optimizer, "DeepSpeed CPU Adam (C++)"); m.def("destroy_adam", &destroy_adam_optimizer, "DeepSpeed CPU Adam destroy (C++)"); diff --git a/csrc/adam/cpu_adam_impl.cpp b/csrc/adam/cpu_adam_impl.cpp index 2adf981dad42..0dc5932d7c75 100644 --- a/csrc/adam/cpu_adam_impl.cpp +++ b/csrc/adam/cpu_adam_impl.cpp @@ -32,9 +32,11 @@ static std::unordered_map> s_optimizers; // C++ interface -template +template void Adam_Optimizer::Step_1(ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, @@ -86,9 +88,11 @@ void Adam_Optimizer::Step_1(ds_params_precision_t* _params, } } -template +template void Adam_Optimizer::Step_4(ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, @@ -155,9 +159,11 @@ int create_adam_optimizer(int optimizer_id, return 0; } -template +template void Adam_Optimizer::Step_8(ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, @@ -180,7 +186,9 @@ void Adam_Optimizer::Step_8(ds_params_precision_t* _params, parallel); } -template +template void step_invoker(std::shared_ptr opt, void* _params, void* grads, @@ -190,7 +198,7 @@ void step_invoker(std::shared_ptr opt, bool parallel) { opt->Step_8((ds_params_precision_t*)(_params), - (ds_params_precision_t*)(grads), + (ds_grad_precision_t*)(grads), (ds_state_precision_t*)(_exp_avg), (ds_state_precision_t*)(_exp_avg_sq), _param_size, @@ -198,17 +206,20 @@ void step_invoker(std::shared_ptr opt, } std::map< - std::tuple, + std::tuple, std::function, void*, void*, void*, void*, size_t, bool)>> invokers; // Fill map with template functions for each type -template +template void create_invoker() { invokers[std::tuple(c10::CppTypeToScalarType(), - c10::CppTypeToScalarType())] = - step_invoker; + c10::CppTypeToScalarType(), + c10::CppTypeToScalarType())] = + step_invoker; } struct InvokerInitializer { InvokerInitializer() @@ -218,6 +229,8 @@ struct InvokerInitializer { create_invoker(); create_invoker(); create_invoker(); + // fp32 weights on bf16 gradients: SuperOffload copies them to the CPU in bf16. + create_invoker(); } } _invoker_initializer; @@ -231,11 +244,13 @@ void invoke(std::shared_ptr opt, { c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype()); c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg.options().dtype()); + c10::ScalarType grad_type = at::typeMetaToScalarType(grads.options().dtype()); - auto it = invokers.find(std::tuple(params_type, state_type)); + auto it = invokers.find(std::tuple(params_type, state_type, grad_type)); if (it == invokers.end()) { throw std::runtime_error("Adam optimizer with param type "s + c10::toString(params_type) + - " and state type "s + c10::toString(state_type) + + ", state type "s + c10::toString(state_type) + " and grad type "s + + c10::toString(grad_type) + " is not supported on current hardware"s); } @@ -604,9 +619,9 @@ class ZenFlowAdam { torch::Tensor& M = grp.exp_avg[now_state]; torch::Tensor& V = grp.exp_avg_sq[now_state]; - auto it = invokers.find(std::tuple(P.scalar_type(), M.scalar_type())); + auto it = invokers.find(std::tuple(P.scalar_type(), M.scalar_type(), G.scalar_type())); TORCH_CHECK(it != invokers.end(), - "ZenFlowAdam: unsupported param/state dtype combination"); + "ZenFlowAdam: unsupported param/state/grad dtype combination"); auto fn = it->second; char* pp = static_cast(P.data_ptr()); @@ -615,13 +630,14 @@ class ZenFlowAdam { char* vp = static_cast(V.data_ptr()); char* sp = grp.stale.defined() ? static_cast(grp.stale.data_ptr()) : nullptr; const size_t pe = P.element_size(); + const size_t ge = G.element_size(); const size_t se = M.element_size(); const size_t numel = P.numel(); pool_->parallel_for(numel, kZenAdamAlign, [=](size_t b, size_t e) { const size_t len = e - b; // parallel=false: each pinned thread runs its slice serially. - fn(opt, pp + b * pe, gp + b * pe, mp + b * se, vp + b * se, len, false); + fn(opt, pp + b * pe, gp + b * ge, mp + b * se, vp + b * se, len, false); if (sp) std::memcpy(sp + b * pe, pp + b * pe, len * pe); }); } diff --git a/csrc/includes/cpu_adam.h b/csrc/includes/cpu_adam.h index abab3ecda999..59d4a246d967 100644 --- a/csrc/includes/cpu_adam.h +++ b/csrc/includes/cpu_adam.h @@ -16,13 +16,15 @@ #include "simd.h" #include "sve.h" -#define STEP(SPAN) \ - template \ - void Step_##SPAN(ds_params_precision_t* _params, \ - ds_params_precision_t* grads, \ - ds_state_precision_t* _exp_avg, \ - ds_state_precision_t* _exp_avg_sq, \ - size_t _param_size, \ +#define STEP(SPAN) \ + template \ + void Step_##SPAN(ds_params_precision_t* _params, \ + ds_grad_precision_t* grads, \ + ds_state_precision_t* _exp_avg, \ + ds_state_precision_t* _exp_avg_sq, \ + size_t _param_size, \ bool parallel = true); class Adam_Optimizer { @@ -47,30 +49,39 @@ class Adam_Optimizer { ~Adam_Optimizer() {} #if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__) - template + template void Step_AVX(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t param_size, bool parallel = true); #endif #if defined(__SVE__) && defined(__ARM_FEATURE_SVE) - template + template void Step_SVE(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t param_size, bool parallel = true); #endif #if defined(__RVV__) && defined(__riscv_v_intrinsic) && __riscv_v_intrinsic >= 12000 - template + template void Step_RVV(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t param_size, @@ -134,10 +145,13 @@ class Adam_Optimizer { }; #if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__) -template +template void Adam_Optimizer::Step_AVX(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, @@ -145,7 +159,8 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, { #if !defined(__AVX512__) && !defined(__NEON__) if (std::is_same_v || - std::is_same_v) { + std::is_same_v || + std::is_same_v) { return; } #endif @@ -225,17 +240,21 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, #endif #if defined(__SVE__) && defined(__ARM_FEATURE_SVE) -template +template void Adam_Optimizer::Step_SVE(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, bool parallel) { if constexpr (std::is_same_v && - std::is_same_v) { + std::is_same_v && + std::is_same_v) { sve_adam_update(_params, grads, _exp_avg, @@ -256,17 +275,21 @@ void Adam_Optimizer::Step_SVE(size_t* rounded_size, #endif #if defined(__RVV__) && defined(__riscv_v_intrinsic) && __riscv_v_intrinsic >= 12000 -template +template void Adam_Optimizer::Step_RVV(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, bool parallel) { if constexpr (std::is_same_v && - std::is_same_v) { + std::is_same_v && + std::is_same_v) { rvv_adam_update(_params, grads, _exp_avg, diff --git a/tests/unit/ops/adam/test_cpu_adam.py b/tests/unit/ops/adam/test_cpu_adam.py index ca925c23c203..280b48c69242 100644 --- a/tests/unit/ops/adam/test_cpu_adam.py +++ b/tests/unit/ops/adam/test_cpu_adam.py @@ -6,6 +6,8 @@ import math import os import sys +import threading +import time import torch import numpy as np import pytest @@ -111,6 +113,38 @@ def test_cpu_adam_strict_state_updates(model_size, adamw_mode, weight_decay, bia torch.testing.assert_close(state['exp_avg_sq'], ref_exp_avg_sq, rtol=3e-5, atol=2e-6) +def test_cpu_adam_step_releases_gil(): + """A CPU Adam step in another thread lets this thread run, as SuperOffload's worker thread needs.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + # Large enough that a step takes several ms, longer than the wakeups of this thread are apart. + param = torch.nn.Parameter(torch.zeros(2**26)) + param.grad = torch.full_like(param, 0.5) + optimizer = DeepSpeedCPUAdam([param]) + optimizer.step() + steps = [] + done = threading.Event() + + def step(): + for _ in range(5): + start = time.perf_counter() + optimizer.step() + steps.append((start, time.perf_counter())) + done.set() + + thread = threading.Thread(target=step) + thread.start() + wakeups = [] + while not done.is_set(): + time.sleep(0.001) + wakeups.append(time.perf_counter()) + thread.join() + # Holding the GIL during a step would keep this thread from waking up until the step is done, so none of its + # wakeups would fall inside a step. Being descheduled now and then only loses a few. + inside = sum(any(a + 1e-3 < t < b - 1e-3 for a, b in steps) for t in wakeups) + assert inside >= 10, f"{inside} wakeups during {len(steps)} steps of {sum(b - a for a, b in steps) * 1e3:.0f} ms" + + def test_cpu_adam_rejects_amsgrad(): from deepspeed.ops.adam import DeepSpeedCPUAdam @@ -123,6 +157,40 @@ def test_cpu_adam_rejects_amsgrad(): assert len(optimizer.param_groups) == 1 +def test_cpu_adam_steps_fp32_params_on_bf16_grads(): + """fp32 params step on bf16 gradients as on the same gradients in fp32, and a gradient dtype that has no kernel + is rejected instead of being read as the params' dtype.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + op = CPUAdamBuilder().load() + opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + lr, beta1, beta2, eps, wd = 1e-3, 0.9, 0.999, 1e-8, 0.01 + n = 100003 # not a multiple of the SIMD block, so the scalar tail runs too + torch.manual_seed(0) + params = torch.randn(n) + grads = [torch.randn(n).to(torch.bfloat16) for _ in range(3)] + results = [] + for grad_dtype in (torch.bfloat16, torch.float): + # A new native optimizer for each, as it carries beta**step over from the steps before. + op.create_adam(opt_id, lr, beta1, beta2, eps, wd, True, False) + p, m, v = params.clone(), torch.zeros(n), torch.zeros(n) + for step, grad in enumerate(grads, 1): + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad.to(grad_dtype), m, v) + results.append((p, m, v)) + op.destroy_adam(opt_id) + for name, a, b in zip(("params", "exp_avg", "exp_avg_sq"), *results): + assert torch.equal(a, b), name + + op.create_adam(opt_id, lr, beta1, beta2, eps, wd, True, False) + try: + with pytest.raises(RuntimeError, match="grad type"): + op.adam_update(opt_id, 1, lr, beta1, beta2, eps, wd, True, params.clone(), grads[0].half(), torch.zeros(n), + torch.zeros(n)) + finally: + op.destroy_adam(opt_id) + + @pytest.mark.parametrize('dtype', [torch.half, torch.bfloat16, torch.float], ids=["fp16", "bf16", "fp32"]) @pytest.mark.parametrize('model_size', [ From f90d9eab9669ff24cf3ef602bc7802f190e2c5d8 Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Thu, 5 Nov 2026 22:24:07 +0000 Subject: [PATCH 07/10] SuperOffload: run CPU Adam in a thread of the training process, copy 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 #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: #8840 pins the master weights there in power-of-two blocks. Signed-off-by: Xinyu Lian --- .../superoffload/superoffload_stage3.py | 537 ++++++++++-------- .../superoffload/superoffload_utils.py | 449 ++++----------- .../unit/v1/pin_memory/test_offload_route.py | 17 - tests/unit/v1/zero/test_invalid_grad_norm.py | 7 +- tests/unit/v1/zero/test_superoffload.py | 109 +++- .../unit/v1/zero/test_zero_partial_offload.py | 6 +- 6 files changed, 492 insertions(+), 633 deletions(-) diff --git a/deepspeed/runtime/superoffload/superoffload_stage3.py b/deepspeed/runtime/superoffload/superoffload_stage3.py index 4633d0e4eb7c..48666660982b 100644 --- a/deepspeed/runtime/superoffload/superoffload_stage3.py +++ b/deepspeed/runtime/superoffload/superoffload_stage3.py @@ -3,15 +3,14 @@ # DeepSpeed Team -import os -import time -from collections import defaultdict -from typing import List +from collections import defaultdict, deque +from concurrent.futures import Future +from typing import List, NamedTuple import torch from deepspeed.ops.adam import DeepSpeedCPUAdam -from deepspeed.runtime.superoffload.superoffload_utils import SuperOffloadCPUOptimizer, TaskKeys, ResultKeys, EventTypes +from deepspeed.runtime.superoffload.superoffload_utils import SuperOffloadCPUAdam from deepspeed.runtime.zero.partition_parameters import Parameter, Tensor from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 from deepspeed.runtime.utils import combine_grad_norm_groups, is_invalid_grad_norm @@ -30,7 +29,10 @@ def _validate_superoffload_accelerator(): def _cpu_adam_settings(optimizer): - """Return (adamw_mode, fp32_optimizer_states) for the DeepSpeedCPUAdam that runs ``optimizer`` in the worker.""" + """Return (adamw_mode, fp32_optimizer_states) for the CPU Adam that steps the CPU sub-groups of ``optimizer``.""" + if any(group.get("amsgrad", False) for group in optimizer.param_groups): + raise ValueError( + "SuperOffload steps the CPU sub-groups with DeepSpeedCPUAdam, which does not support AMSGrad.") if isinstance(optimizer, DeepSpeedCPUAdam): return optimizer.adam_w_mode, optimizer.fp32_optimizer_states if isinstance(optimizer, torch.optim.AdamW): @@ -41,6 +43,14 @@ def _cpu_adam_settings(optimizer): f"AdamW, but the optimizer is {type(optimizer).__name__}.") +class _CPUStep(NamedTuple): + """A CPU step submitted during backward, and what a rollback of it needs.""" + # Its result tells whether the sub-group stepped: one whose gradients were not finite is left to step(). + stepped: Future + hyperparams: dict + grad_scale: float + + class SuperOffloadOptimizer_Stage3(DeepSpeedZeroOptimizer_Stage3): def __init__( @@ -55,92 +65,114 @@ def __init__( _validate_superoffload_accelerator() adamw_mode, fp32_optimizer_states = _cpu_adam_settings(init_optimizer) - self.sub_group_to_param_num = {} # Per sub-group, the ds_ids of the params whose gradient for this step is in the fp32 gradient buffer. self.sub_group_grads_ready = defaultdict(set) - # Sub-groups stepped during backward, mapped to the (hyperparameters, grad_scale) of that step for a rollback. + # Sub-groups whose CPU step was submitted during backward, until step() is done or the step is rolled back. self._submitted_cpu_sub_groups = {} - self.async_cpuadam_num = 0 - self.max_grad_numel = 0 + # The submitted sub-groups whose step has not been collected yet, in the order the CPU Adam thread runs them. + self._uncollected_cpu_steps = deque() + # Host addresses of the CPU sub-groups' fp32 weights and gradients registered with CUDA. + self._registered_host_addresses = [] + # Sub-groups stepped during backward whose new weights are converted into a GPU staging buffer, mapped to the + # event of that conversion. + self._staged_lp_params = {} + self._lp_staging_buffers = {} + # Per staging buffer, the event after which step() has copied it into the params. + self._lp_staging_free_events = {} + self._lp_staging_stream = None + # Stream of the gradient copies to the CPU, so that they overlap the backward compute. + self._grad_copy_stream = None super().__init__(module, init_optimizer, param_names, timers, ds_config, **kwargs) - optimizer_configs = [] - for pg in self.optimizer.param_groups: - optimizer_configs.append({ - "lr": pg["lr"], - "betas": pg["betas"], - "eps": pg["eps"], - "weight_decay": pg["weight_decay"], - "bias_correction": pg.get("bias_correction", True), - "amsgrad": pg["amsgrad"], - }) + self.sub_group_to_param_num = [len(group) for group in self.fp16_groups] # A step taken during backward is undone with adam_rollback, which supports only fp32 AdamW with bias # correction. With other settings, step() runs the CPU sub-groups after the overflow and clipping checks. + bias_correction = all(group.get("bias_correction", True) for group in self.optimizer.param_groups) self.speculative_cpu_step = (adamw_mode and fp32_optimizer_states - and self.master_weights_and_grads_dtype == torch.float32 - and all(config["bias_correction"] for config in optimizer_configs)) - cpuadam_cores_perc = kwargs.get("cpuadam_cores_perc", 0.8) - self.superoffload_cpu_optimizer = SuperOffloadCPUOptimizer(optimizer_config=optimizer_configs, - cpuadam_cores_perc=cpuadam_cores_perc, - max_grad_numel=self.max_grad_numel, - pin_memory=self.offload_optimizer_pin_memory, - adamw_mode=adamw_mode, - fp32_optimizer_states=fp32_optimizer_states, - grad_dtype=self.master_weights_and_grads_dtype) - self._check_shm_size(fp32_optimizer_states) - - def _check_shm_size(self, fp32_optimizer_states): - """Warn when /dev/shm is too small for the tensors that SuperOffload shares with its worker process.""" - try: - stat = os.statvfs("/dev/shm") - except OSError: + and self.master_weights_and_grads_dtype == torch.float32 and bias_correction) + self.cpu_adam = SuperOffloadCPUAdam(self.optimizer.param_groups[0], adamw_mode, fp32_optimizer_states, + kwargs.get("cpuadam_cores_perc", 0.8)) + # With bf16 gradients accumulated into fp32 master weights in one micro-step, the gradients of the CPU + # sub-groups go to the CPU in bf16, half the bytes of fp32, and the worker promotes them in its kernel. + self._half_cpu_grads = {} + if (self.dtype == torch.bfloat16 and self.gradient_accumulation_dtype == torch.bfloat16 + and self.master_weights_and_grads_dtype == torch.float32 and fp32_optimizer_states + and self.gradient_accumulation_steps == 1): + for sub_group_id, device in self.subgroup_to_device.items(): + if device == 'cpu' and not self._swappable_optimizer_subgroup(sub_group_id): + self._half_cpu_grads[sub_group_id] = torch.zeros( + self.fp32_partitioned_groups_flat[sub_group_id].numel(), dtype=torch.bfloat16) + self._pin_cpu_buffers() + self._grad_copy_stream = get_accelerator().Stream() + if self.speculative_cpu_step and self._cpu_sub_group_weights_pinned(): + # Default priority is the lowest, so the conversions never delay the backward kernels. + self._lp_staging_stream = get_accelerator().Stream() + + def _pin_cpu_buffers(self): + """Keep the CPU sub-groups' fp32 weights and gradients (and bf16 gradients) page-locked. + + offload_optimizer.pin_memory pins the gradients but not the fp32 weights. The gradient copies of backward and + the weight copies of step() would then run from pageable memory, at a fraction of the bandwidth and blocking + the thread that issues them. Registering their memory with the device runtime keeps the copies asynchronous. + """ + if not self.offload_optimizer_pin_memory: return - numel = sum(self.fp32_partitioned_groups_flat[i].numel() for i, device in self.subgroup_to_device.items() - if device == 'cpu') - weight_size = torch.empty(0, dtype=self.master_weights_and_grads_dtype).element_size() - state_size = 4 if fp32_optimizer_states else weight_size - # The fp32 weights and gradients move there at the first step, the Adam state at the first checkpoint save - # or load. - needed = numel * 2 * (weight_size + state_size) - local_ranks = int(os.environ.get("LOCAL_WORLD_SIZE", os.environ.get("LOCAL_SIZE", 1))) - size = stat.f_blocks * stat.f_frsize - if size < needed * local_ranks: - logger.warning( - f"SuperOffload shares the fp32 weights, gradients and Adam state of the CPU sub-groups with " - f"its worker process through /dev/shm: {needed / 2**30:.1f} GiB on this rank, " - f"{local_ranks} rank(s) on this node, but /dev/shm has {size / 2**30:.1f} GiB. A checkpoint " - f"save or load may fail; increase the size of /dev/shm.") - - def _create_fp16_sub_groups(self, params_group): - - params_group_numel = sum([param.partition_numel() for param in params_group]) - sub_group_size = self.sub_group_size - - if sub_group_size is None or sub_group_size >= params_group_numel: - global_idx = len(self.sub_group_to_param_num) - self.sub_group_to_param_num[global_idx] = len(params_group) - self.max_grad_numel = max(self.max_grad_numel, params_group_numel) - return [params_group] - - sub_groups = [] - sub_group = [] - local_sub_group_size = 0 - - for param in params_group: - sub_group.append(param) - local_sub_group_size += param.partition_numel() - - if local_sub_group_size >= sub_group_size or id(param) == id(params_group[-1]): - self.max_grad_numel = max(self.max_grad_numel, local_sub_group_size) - sub_groups.append(sub_group) - global_idx = len(self.sub_group_to_param_num) - self.sub_group_to_param_num[global_idx] = len(sub_group) - - sub_group = [] - local_sub_group_size = 0 - - return sub_groups + accelerator = get_accelerator() + for sub_group_id, device in self.subgroup_to_device.items(): + # Swapped sub-groups share one gradient buffer and do not keep their weights in memory. + if device != 'cpu' or self._swappable_optimizer_subgroup(sub_group_id): + continue + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + for tensor in (fp32_param.data, fp32_param.grad, self._half_cpu_grads.get(sub_group_id)): + if tensor is None or tensor.is_pinned(): + continue + storage = tensor.untyped_storage() + if storage.data_ptr() in self._registered_host_addresses: + continue + try: + registered = accelerator.register_host_memory(storage.data_ptr(), storage.nbytes()) + except Exception as e: + registered = False + logger.warning(f"SuperOffload could not page-lock its CPU buffers: {e}") + if not registered: + return + self._registered_host_addresses.append(storage.data_ptr()) + + def _cpu_sub_group_weights_pinned(self): + """Whether the fp32 weights of the CPU sub-groups are page-locked, which the GPU staging of step() needs.""" + weights = [ + self.fp32_partitioned_groups_flat[sub_group_id].data + for sub_group_id, device in self.subgroup_to_device.items() + if device == 'cpu' and not self._swappable_optimizer_subgroup(sub_group_id) + ] + return bool(weights) and all(weight.is_pinned() for weight in weights) + + def _worker_grad(self, sub_group_id): + """The gradient of a sub-group that the CPU worker steps on: the bf16 one if it has one, else the fp32 one.""" + half_grad = self._half_cpu_grads.get(sub_group_id) + return half_grad if half_grad is not None else self.fp32_partitioned_groups_flat[sub_group_id].grad.data + + def _cpu_step_grad(self, sub_group_id, total_norm): + """The gradient for step() to step a CPU sub-group with that did not step during backward, unscaled and + clipped. A bf16 gradient that needs neither goes as it is; otherwise the fp32 gradient takes the scaling.""" + half_grad = self._half_cpu_grads.get(sub_group_id) + if half_grad is not None: + if self.loss_scale == 1 and not self.check_clip_grads(total_norm): + return half_grad + self.fp32_partitioned_groups_flat[sub_group_id].grad.copy_(half_grad) + self.unscale_and_clip_grads(sub_group_id, total_norm) + return self.fp32_partitioned_groups_flat[sub_group_id].grad.data + + def destroy(self): + self.cpu_adam.close() + super().destroy() + + def _unpin_offload_buffers(self): + for address in self._registered_host_addresses: + get_accelerator().unregister_host_memory(address) + self._registered_host_addresses = [] + super()._unpin_offload_buffers() def _optimizer_step(self, sub_group_id): param_group_id = self.sub_group_to_group_id[sub_group_id] @@ -163,8 +195,13 @@ def step_with_gradscaler(optimizer): @instrument_w_nvtx def _reassign_or_swap_out_partitioned_parameters(self, sub_group_id): if self.subgroup_to_device[sub_group_id] == 'cpu': - self.fp16_partitioned_groups_flat[sub_group_id].data.copy_( - self.fp32_partitioned_groups_flat[sub_group_id].data) + lp_flat = self.fp16_partitioned_groups_flat[sub_group_id] + fp32_flat = self.fp32_partitioned_groups_flat[sub_group_id].data + if fp32_flat.is_pinned() and get_accelerator().on_accelerator(lp_flat): + # copy_ from a CPU tensor of another dtype converts on the CPU into a pageable temporary and copies + # that. Copying the pinned fp32 weights and converting on the GPU keeps the CPU free and the copy fast. + fp32_flat = fp32_flat.to(lp_flat.device, non_blocking=True) + lp_flat.data.copy_(fp32_flat) self._unflatten_partitioned_parameters(sub_group_id) return @@ -176,9 +213,79 @@ def _reassign_or_swap_out_partitioned_parameters(self, sub_group_id): self._partitioned_params_swap_out(sub_group_id) @instrument_w_nvtx - def _reassign_or_swap_out_partitioned_parameters_async(self, sub_group_id, updated_param): - """Asynchronously update partitioned parameters with optimized values.""" - self.fp32_partitioned_groups_flat[sub_group_id].data.copy_(updated_param, non_blocking=True) + def _stage_lp_params(self, sub_group_id): + """Convert the weights of a CPU step taken during backward to the low-precision dtype on the GPU. + + Backward may still gather the current params, and step() changes them only after the overflow and clipping + checks, so the conversion goes to a staging buffer that step() copies in. + """ + stream = self._lp_staging_stream + lp_flat = self.fp16_partitioned_groups_flat[sub_group_id] + fp32_flat = self.fp32_partitioned_groups_flat[sub_group_id].data + if stream is None or lp_flat is None or not get_accelerator().on_accelerator(lp_flat): + return + if not fp32_flat.is_pinned(): + return + staging = self._lp_staging_buffers.get(sub_group_id) + if staging is None: + # Allocate it on the stream that writes it. The caching allocator hands the current stream blocks freed + # while kernels queued on that stream may still use them, which another stream does not wait for. + with get_accelerator().stream(stream): + staging = torch.empty_like(lp_flat) + self._lp_staging_buffers[sub_group_id] = staging + free_event = self._lp_staging_free_events.pop(sub_group_id, None) + if free_event is not None: + # The previous step() may still be copying this buffer into the params. + stream.wait_event(free_event) + with get_accelerator().stream(stream): + staging.copy_(fp32_flat.to(staging.device, non_blocking=True)) + event = get_accelerator().Event() + event.record(stream) + self._staged_lp_params[sub_group_id] = event + + def _drop_staged_lp_params(self, sub_group_id): + """Discard a staged conversion before the worker changes the sub-group's weights, once it has read them.""" + event = self._staged_lp_params.pop(sub_group_id, None) + if event is not None: + event.synchronize() + + @instrument_w_nvtx + def _apply_staged_lp_params(self, sub_group_id): + event = self._staged_lp_params.pop(sub_group_id) + stream = get_accelerator().current_stream() + stream.wait_event(event) + self.fp16_partitioned_groups_flat[sub_group_id].data.copy_(self._lp_staging_buffers[sub_group_id]) + free_event = get_accelerator().Event() + free_event.record(stream) + self._lp_staging_free_events[sub_group_id] = free_event + self._unflatten_partitioned_parameters(sub_group_id) + + def _collect_cpu_steps(self, wait=False): + """Stage the new weights of the CPU steps submitted during backward that have finished. With wait, wait for + all of them, and raise what a step raised. + + The CPU Adam thread runs the steps in the order they were submitted, so they finish in that order. + """ + while self._uncollected_cpu_steps: + sub_group_id = self._uncollected_cpu_steps[0] + stepped = self._submitted_cpu_sub_groups[sub_group_id].stepped + if not wait and not stepped.done(): + return + self._uncollected_cpu_steps.popleft() + if stepped.result(): + self._stage_lp_params(sub_group_id) + else: + # Its gradients were not finite, so step() runs it after the overflow check. + del self._submitted_cpu_sub_groups[sub_group_id] + + def _copies_grad_directly(self, sub_group_id): + """Whether a gradient of the sub-group goes straight into the master-dtype gradient the worker reads. + + That needs a CPU sub-group and the only micro-step of the step, so that no earlier micro-step left a gradient + in the reduce buffer. fp16 keeps the reduce buffer, which its overflow check scans. + """ + return (self.subgroup_to_device[sub_group_id] == 'cpu' and self.dtype != torch.float16 + and self.is_gradient_accumulation_boundary() and self.micro_step_id == 0) @instrument_w_nvtx def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None: @@ -195,7 +302,24 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L # segments use. first_grad = param.ds_id not in self._offload_boundary_param_ids self._offload_boundary_param_ids.add(param.ds_id) - if first_grad: + if self._copies_grad_directly(i): + if first_grad: + # The gradient goes to the CPU once, into the gradient the worker reads, instead of into the + # reduce buffer, back to the GPU for the norm, and out again. It is rounded to the accumulation + # dtype as the reduce buffer would round it, and goes in that dtype if the worker reads bf16. + grad_buffer = grad_partition.to(self.gradient_accumulation_dtype) + if i not in self._half_cpu_grads: + grad_buffer = grad_buffer.to(self.master_weights_and_grads_dtype) + else: + # A second gradient in this step: the first one is only in the gradient the worker reads, which + # the copy stream may still be writing. The sum is taken in the accumulation dtype, as in the + # reduce buffer; the first gradient was rounded to it, so converting it back is exact. + cpu_grad_tensor = self._worker_grad(i).narrow(0, dest_offset, grad_partition.numel()) + get_accelerator().current_stream().wait_stream(self._grad_copy_stream) + grad_buffer = cpu_grad_tensor.to(grad_partition.device, non_blocking=True) + grad_buffer = grad_buffer.to(self.gradient_accumulation_dtype) + grad_buffer.add_(grad_partition.to(grad_buffer.dtype).view(grad_buffer.shape)) + elif first_grad: grad_buffer.copy_(grad_partition, non_blocking=True) grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) elif get_accelerator().on_accelerator(grad_buffer): @@ -209,45 +333,61 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L if self.is_gradient_accumulation_boundary(): if i in self._submitted_cpu_sub_groups: # Another gradient for a sub-group that has already stepped: undo that step while the fp32 - # buffer still holds the gradient it used. - self._wait_for_async_operations() - self._rollback_cpu_sub_group(i) + # buffer still holds the gradient it used. One whose gradients were not finite did not step. + self._collect_cpu_steps(wait=True) + if i in self._submitted_cpu_sub_groups: + self._rollback_cpu_sub_group(i) self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) - fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow( - 0, dest_offset, grad_buffer.numel()) - fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype), non_blocking=True) + cpu_grad_tensor = self._worker_grad(i).narrow(0, dest_offset, grad_buffer.numel()) + grad_to_copy = grad_buffer.to(dtype=cpu_grad_tensor.dtype) + copy_stream = self._grad_copy_stream + copy_stream.wait_stream(get_accelerator().current_stream()) + with get_accelerator().stream(copy_stream): + cpu_grad_tensor.copy_(grad_to_copy, non_blocking=True) + # The copy reads the gradient later, so keep the allocator from handing its memory out before then. + grad_to_copy.record_stream(copy_stream) self.sub_group_grads_ready[i].add(param.ds_id) if len(self.sub_group_grads_ready[i]) == self.sub_group_to_param_num[i]: completed_sub_groups.append(i) if self.is_gradient_accumulation_boundary() and completed_sub_groups: - get_accelerator().current_stream().synchronize() + steps = [] + finite_on_device = [] # A param that gets a second gradient adds its sub-group again. for i in dict.fromkeys(completed_sub_groups): param_group_id = self.sub_group_to_group_id[i] if (self.subgroup_to_device[i] == 'cpu' and self.speculative_cpu_step - and self._betas_nonzero(param_group_id) and self._grads_finite(i)): - fp32_param = self.fp32_partitioned_groups_flat[i] + and self._betas_nonzero(param_group_id)): hyperparams = self._hyperparams(param_group_id) # The gradients still carry the loss scale, which step() would remove. grad_scale = 1.0 / self.loss_scale - self.superoffload_cpu_optimizer.async_step(param_group_id, - i, - fp32_param.data, - fp32_param.grad.data, - hyperparams=hyperparams, - grad_scale=grad_scale) - self._submitted_cpu_sub_groups[i] = (hyperparams, grad_scale) - self.async_cpuadam_num += 1 - - result = self.superoffload_cpu_optimizer.get_result() - if result is not None: - self._reassign_or_swap_out_partitioned_parameters_async(result[TaskKeys.SUB_GROUP_ID], - result[ResultKeys.UPDATED_PARAM]) - self.async_cpuadam_num -= 1 + self._drop_staged_lp_params(i) + steps.append((i, hyperparams, grad_scale)) + finite_on_device.append(self._grads_finite(i)) + if steps: + # The CPU Adam thread reads the gradients that the copy stream writes, and waits for them itself, so + # that this thread goes on queueing the backward kernels. + copy_stream = self._grad_copy_stream + copy_stream.wait_stream(get_accelerator().current_stream()) + with get_accelerator().stream(copy_stream): + # Behind the gradient copies: a copy to the CPU on the compute stream would wait for them in the + # copy engine, and hold up the kernels queued after it. + finite = [flag.to("cpu", non_blocking=True) for flag in finite_on_device] + copies_done = get_accelerator().Event() + copies_done.record(copy_stream) + for flag in finite_on_device: + flag.record_stream(copy_stream) + for (i, hyperparams, grad_scale), grads_finite in zip(steps, finite): + stepped = self.cpu_adam.submit(self._step_when_copied, i, copies_done, grads_finite, hyperparams, + grad_scale) + self._submitted_cpu_sub_groups[i] = _CPUStep(stepped, hyperparams, grad_scale) + self._uncollected_cpu_steps.append(i) + + if self.is_gradient_accumulation_boundary(): + self._collect_cpu_steps() for param in params_to_release: if not get_accelerator().is_synchronized_device(): @@ -260,7 +400,9 @@ def step(self, closure=None): """ Not supporting closure. """ - self._wait_for_async_operations() + # Sub-groups that did not complete in backward may still have gradient copies in flight. + self._grad_copy_stream.synchronize() + self._collect_cpu_steps(wait=True) self._complete_grads() self._pre_step() @@ -293,20 +435,20 @@ def step(self, closure=None): for sub_group_id, group in enumerate(self.fp16_groups): self._prepare_sub_group(sub_group_id, timer_names) - self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) - + # A sub-group stepped during backward had the loss scale removed in the worker, and one that needed + # clipping was rolled back above, so its gradients are not read again. if self.subgroup_to_device[sub_group_id] != 'cpu': + self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) self._optimizer_step(sub_group_id) elif sub_group_id not in self._submitted_cpu_sub_groups: param_group_id = self.sub_group_to_group_id[sub_group_id] - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - self._sync_cpu_optimizer_step(param_group_id, - sub_group_id, - fp32_param.data, - fp32_param.grad.data, - hyperparams=self._hyperparams(param_group_id)) - - self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + self._cpu_step_and_wait(sub_group_id, self._cpu_step_grad(sub_group_id, scaled_global_grad_norm), + self._hyperparams(param_group_id)) + + if sub_group_id in self._staged_lp_params: + self._apply_staged_lp_params(sub_group_id) + else: + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) self._release_sub_group(sub_group_id, timer_names) self.timers(OPTIMIZER_STEP_TIMER).stop() @@ -314,10 +456,12 @@ def step(self, closure=None): def _post_step(self, timer_names): super()._post_step(timer_names) + for sub_group_id in list(self._staged_lp_params): + self._drop_staged_lp_params(sub_group_id) self._submitted_cpu_sub_groups.clear() def _complete_grads(self): - """Fill the fp32 gradient of the params that got no gradient in the boundary micro-step. + """Fill the gradient the worker reads for the params that got no gradient in the boundary micro-step. A param used only in earlier micro-steps has its gradient in the reduce buffer. A param not used in this step gets a zero gradient, as set_none_gradients_to_zero gives it in ZeRO-1/2; zeroing its reduce buffer also keeps @@ -340,17 +484,18 @@ def _complete_grads(self): if not get_accelerator().on_accelerator(grad_buffer): grad_buffer = grad_buffer.to(get_accelerator().current_device_name(), non_blocking=True) self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) - fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow(0, dest_offset, num_elements) - fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype), non_blocking=True) + cpu_grad_tensor = self._worker_grad(i).narrow(0, dest_offset, num_elements) + cpu_grad_tensor.copy_(grad_buffer.to(dtype=cpu_grad_tensor.dtype), non_blocking=True) filled = True if filled: get_accelerator().current_stream().synchronize() self.sub_group_grads_ready.clear() def _grads_finite(self, sub_group_id): - """Whether the sub-group's gradients are finite. A step on inf or nan cannot be rolled back.""" + """Whether the sub-group's gradients are finite, as a device tensor, so that this thread does not wait for the + norms. A step on inf or nan cannot be rolled back.""" norms = [self.norm_for_param_grads[self.get_param_id(param)] for param in self.fp16_groups[sub_group_id]] - return bool(torch.isfinite(torch.stack(norms)).all()) + return torch.isfinite(torch.stack(norms)).all() def _betas_nonzero(self, param_group_id): """Whether a step of the param group can be rolled back: the rollback divides the moments by beta1 and beta2. @@ -360,91 +505,35 @@ def _betas_nonzero(self, param_group_id): def _hyperparams(self, param_group_id): """The hyperparameters of a param group, which LR, momentum or weight decay schedules change.""" group = self.optimizer.param_groups[param_group_id] - return {key: group[key] for key in ("lr", "betas", "eps", "weight_decay")} + hyperparams = {key: group[key] for key in ("lr", "betas", "eps", "weight_decay")} + hyperparams["bias_correction"] = group.get("bias_correction", True) + return hyperparams def _loss_scale_changed(self): grad_scale = 1.0 / self.loss_scale - return any(scale != grad_scale for _, scale in self._submitted_cpu_sub_groups.values()) + return any(step.grad_scale != grad_scale for step in self._submitted_cpu_sub_groups.values()) - def _wait_for_async_operations(self, timeout_seconds=60): - """Wait for all pending asynchronous CPU optimizer operations to complete with timeout error. + def _step_when_copied(self, sub_group_id, copies_done, grads_finite, hyperparams, grad_scale): + """In the CPU Adam thread: step a sub-group once its gradients are on the CPU, and tell whether it did. - Args: - timeout_seconds (int): Maximum time to wait before throwing an error. Default is 60 seconds. + A step on inf or nan cannot be rolled back, so step() runs those sub-groups. """ - if self.async_cpuadam_num > 0: - logger.info(f"[INFO] {self.async_cpuadam_num} asynchronous CPU optimizer operations pending...") - if self.async_cpuadam_num == 0: - return - - start_time = time.time() - initial_pending_ops = self.async_cpuadam_num - - while self.async_cpuadam_num > 0: - result = self.superoffload_cpu_optimizer.get_result() - if result is None: - current_time = time.time() - elapsed_time = current_time - start_time - - # Throw error if we've been waiting longer than the timeout - if elapsed_time >= timeout_seconds: - raise RuntimeError( - f"SuperOffload CPU optimizer timeout after {elapsed_time:.1f} seconds. " - f"Still waiting for {self.async_cpuadam_num}/{initial_pending_ops} async operations to complete. " - f"This indicates a deadlock or critical performance issue in the CPU optimizer.") - - time.sleep(0.001) # 1ms sleep - continue + copies_done.synchronize() + if not grads_finite.item(): + return False + self._cpu_adam_step(sub_group_id, self._worker_grad(sub_group_id), hyperparams, grad_scale) + return True - self._reassign_or_swap_out_partitioned_parameters_async(result[TaskKeys.SUB_GROUP_ID], - result[ResultKeys.UPDATED_PARAM]) - self.async_cpuadam_num -= 1 + def _cpu_adam_step(self, sub_group_id, grad, hyperparams, grad_scale=1.0, rollback=False): + """In the CPU Adam thread: step a CPU sub-group, or undo its last step.""" + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + update = self.cpu_adam.rollback if rollback else self.cpu_adam.step + # The Adam state stays where ZeRO-3 CPU offload keeps it, so that checkpoints save and load it as they are. + update(self.optimizer.state[fp32_param], fp32_param.data, grad, hyperparams, grad_scale) - def _wait_for_single_async_result(self, event_type: str, timeout_seconds=60): - """Wait for a single asynchronous CPU-Adam optimizer operation with timeout. - - Args: - event_type (str): Type of operation expected ('adam_step' or 'rollback'). - timeout_seconds (int): Maximum time to wait before throwing an error. Default is 60 seconds. - """ - start_time = time.time() - - while True: - result = self.superoffload_cpu_optimizer.get_result(expected_event_type=event_type) - if result is not None: - self._reassign_or_swap_out_partitioned_parameters_async(result[TaskKeys.SUB_GROUP_ID], - result[ResultKeys.UPDATED_PARAM]) - break - - current_time = time.time() - elapsed_time = current_time - start_time - - # Throw error if we've been waiting longer than the timeout - if elapsed_time >= timeout_seconds: - raise RuntimeError(f"SuperOffload CPU optimizer timeout after {elapsed_time:.1f} seconds. " - f"This indicates a deadlock or critical performance issue in the CPU optimizer.") - - time.sleep(0.001) # 1ms sleep - - def _sync_cpu_optimizer_step(self, - param_group_id: int, - sub_group_id: int, - fp32_param_data, - fp32_grad_data, - rollback: bool = False, - hyperparams: dict = None, - timeout_seconds: int = 60, - grad_scale: float = 1.0): - event_type = EventTypes.ROLLBACK if rollback else EventTypes.ADAM_STEP - self.superoffload_cpu_optimizer.async_step(param_group_id, - sub_group_id, - fp32_param_data, - fp32_grad_data, - rollback=rollback, - hyperparams=hyperparams, - grad_scale=grad_scale) - # Wait for completion - self._wait_for_single_async_result(event_type, timeout_seconds) + def _cpu_step_and_wait(self, sub_group_id, grad, hyperparams, grad_scale=1.0, rollback=False): + self._drop_staged_lp_params(sub_group_id) + self.cpu_adam.submit(self._cpu_adam_step, sub_group_id, grad, hyperparams, grad_scale, rollback).result() def _rollback_submitted_cpu_sub_groups(self): """Undo the CPU steps taken during backward, for a skipped step or before redoing them with clipping.""" @@ -452,16 +541,13 @@ def _rollback_submitted_cpu_sub_groups(self): self._rollback_cpu_sub_group(sub_group_id) def _rollback_cpu_sub_group(self, sub_group_id): - hyperparams, grad_scale = self._submitted_cpu_sub_groups.pop(sub_group_id) - param_group_id = self.sub_group_to_group_id[sub_group_id] - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - self._sync_cpu_optimizer_step(param_group_id, - sub_group_id, - fp32_param.data, - fp32_param.grad.data, - rollback=True, - hyperparams=hyperparams, - grad_scale=grad_scale) + step = self._submitted_cpu_sub_groups.pop(sub_group_id) + # The rollback reads the gradient that the step read. + self._cpu_step_and_wait(sub_group_id, + self._worker_grad(sub_group_id), + step.hyperparams, + step.grad_scale, + rollback=True) @instrument_w_nvtx def check_clip_grads(self, total_norm): @@ -471,20 +557,6 @@ def check_clip_grads(self, total_norm): clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad return bool(clip > 1.0) - def state_dict(self): - # The worker process holds the Adam state of the CPU sub-groups. Save it where ZeRO-3 CPU offload keeps it. - self._wait_for_async_operations() - cpu_params = [] - for sub_group_id, state in self.superoffload_cpu_optimizer.get_state().items(): - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - self.optimizer.state[fp32_param] = state - cpu_params.append(fp32_param) - try: - return super().state_dict() - finally: - for fp32_param in cpu_params: - del self.optimizer.state[fp32_param] - def load_state_dict(self, state_dict_list, load_optimizer_states=True, @@ -492,17 +564,10 @@ def load_state_dict(self, checkpoint_folder=None, load_serial=None, param_shapes=None): + get_accelerator().current_stream().synchronize() super().load_state_dict(state_dict_list, load_optimizer_states=load_optimizer_states, load_from_fp32_weights=load_from_fp32_weights, checkpoint_folder=checkpoint_folder, load_serial=load_serial, param_shapes=param_shapes) - if not load_optimizer_states: - return - # Move the loaded Adam state of the CPU sub-groups to the worker process, which steps them. - cpu_state = {} - for sub_group_id, fp32_param in enumerate(self.fp32_partitioned_groups_flat): - if self.subgroup_to_device[sub_group_id] == 'cpu' and fp32_param in self.optimizer.state: - cpu_state[sub_group_id] = self.optimizer.state.pop(fp32_param) - self.superoffload_cpu_optimizer.set_state(cpu_state) diff --git a/deepspeed/runtime/superoffload/superoffload_utils.py b/deepspeed/runtime/superoffload/superoffload_utils.py index a6e5dff072d6..d036ab7dd495 100644 --- a/deepspeed/runtime/superoffload/superoffload_utils.py +++ b/deepspeed/runtime/superoffload/superoffload_utils.py @@ -3,378 +3,119 @@ # DeepSpeed Team """ -SuperOffload utilities for 1) running CPU optimizers in separate processes. - +SuperOffload's CPU Adam, which steps the CPU sub-groups in a thread of the training process. """ -import time -from typing import Dict, Optional, Any +import os +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Any, Callable, Dict, List + import torch -import torch.multiprocessing as mp -import psutil from deepspeed.ops.adam import DeepSpeedCPUAdam -from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import CPUAdamBuilder from deepspeed.utils import logger +from deepspeed.utils.logging import should_log_le -class TaskKeys: - PARAM_DATA = "param_data" - PARAM_GRAD = "param_grad" - PARAM_GROUP_ID = "param_group_id" - SUB_GROUP_ID = "sub_group_id" - ROLLBACK = "rollback" - HYPERPARAMS = "hyperparams" - GRAD_SCALE = "grad_scale" - COMMAND = "command" - STATE = "state" - - -class ResultKeys: - UPDATED_PARAM = "updated_param" - EVENT_TYPE = "event_type" - STATE = "state" - - -class EventTypes: - ADAM_STEP = "adam_step" - ROLLBACK = "rollback" - GET_STATE = "get_state" - SET_STATE = "set_state" - - -def _allocate_worker_grad_buffer(max_grad_numel: int, - pin_memory: bool, - dtype: torch.dtype = torch.float32) -> torch.Tensor: - # Scratch destination for the GPU->CPU grad copy. Pinning is optional: it - # speeds DMA but is not required for correctness, so honor - # offload_optimizer.pin_memory rather than always page-locking. - # dtype is the master weights dtype: the optimizer requires the gradient to match the parameter. - if pin_memory: - return get_accelerator().pin_empty(max_grad_numel, dtype=dtype, device='cpu') - return torch.empty(max_grad_numel, dtype=dtype, device='cpu') - - -def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, - result_queue: mp.SimpleQueue, - optimizer_config: Dict[str, Any], - max_grad_numel: int, - pin_memory: bool = True, - adamw_mode: bool = True, - fp32_optimizer_states: bool = True, - grad_dtype: torch.dtype = torch.float32) -> None: - """ - This function runs in a separate process and continuously processes optimization - tasks from the parameter queue. It creates a DeepSpeedCPUAdam optimizer and - applies optimization steps to parameters received from the main process. - - Args: - param_queue: Queue for receiving optimization tasks - result_queue: Queue for sending back optimization results - optimizer_config: Configuration dictionary for the optimizer containing - lr, betas, eps, weight_decay, bias_correction, and amsgrad parameters - max_grad_numel: Maximum number of elements expected in gradient tensors - pin_memory: Whether to page-lock the reusable grad scratch buffer - adamw_mode: Whether the optimizer is AdamW (decoupled weight decay) or Adam - fp32_optimizer_states: Whether to keep the optimizer states in fp32 for low-precision parameters - grad_dtype: dtype of the master weights and their gradients - """ - cpu_tensor = torch.randn(1, device="cpu") - cpu_param = torch.nn.Parameter(cpu_tensor) - - try: - if isinstance(optimizer_config, list): - pg_configs = optimizer_config - else: - pg_configs = [optimizer_config] - - first_cfg = pg_configs[0] - optimizer = DeepSpeedCPUAdam([cpu_param], - lr=first_cfg["lr"], - bias_correction=first_cfg["bias_correction"], - betas=first_cfg["betas"], - eps=first_cfg["eps"], - weight_decay=first_cfg["weight_decay"], - amsgrad=first_cfg["amsgrad"], - adamw_mode=adamw_mode, - fp32_optimizer_states=fp32_optimizer_states) - for cfg in pg_configs[1:]: - dummy = torch.nn.Parameter(torch.randn(1, device="cpu")) - optimizer.add_param_group({ - "params": [dummy], - "lr": cfg["lr"], - "betas": cfg["betas"], - "eps": cfg["eps"], - "weight_decay": cfg["weight_decay"], - "bias_correction": cfg["bias_correction"], - "amsgrad": cfg["amsgrad"], - }) - except KeyError as e: - error_msg = f"Missing required optimizer config key: {e}" - logger.error(error_msg) - result_queue.put({"error": error_msg}) +def _bind_thread(cores: List[int]) -> None: + """Run the calling thread on these cores; the threads it starts later, such as its OpenMP threads, inherit them.""" + if not cores: return + try: + os.sched_setaffinity(0, cores) + except (AttributeError, OSError) as e: + logger.debug(f"Could not set the CPU affinity of a SuperOffload thread: {e}") - # Pre-allocate reusable host buffer for gradients - grad_buffer = _allocate_worker_grad_buffer(max_grad_numel, pin_memory, grad_dtype) - - while True: - try: - task = param_queue.get() - - if task is None: - logger.debug("Received termination signal, shutting down worker") - break - - # The Adam state lives in this process, so checkpoints get and set it here. Its tensors move to shared - # memory when they are sent, without a copy that stays alive. - command = task.get(TaskKeys.COMMAND) - if command == EventTypes.GET_STATE: - state = {sub_group_id: state for sub_group_id, state in optimizer.state.items() if state} - result_queue.put({ResultKeys.EVENT_TYPE: EventTypes.GET_STATE, ResultKeys.STATE: state}) - continue - if command == EventTypes.SET_STATE: - optimizer.state.clear() - optimizer.state.update(task[TaskKeys.STATE]) - result_queue.put({ResultKeys.EVENT_TYPE: EventTypes.SET_STATE}) - continue - - param_data = task[TaskKeys.PARAM_DATA] - param_grad = task[TaskKeys.PARAM_GRAD] - param_group_id = task[TaskKeys.PARAM_GROUP_ID] - sub_group_id = task[TaskKeys.SUB_GROUP_ID] - rollback = task.get(TaskKeys.ROLLBACK, False) - hyperparams = task.get(TaskKeys.HYPERPARAMS, None) - grad_scale = task.get(TaskKeys.GRAD_SCALE, 1.0) - - logger.debug(f"Processing param_group_id: {param_group_id}, sub_group_id: {sub_group_id}") - - del task[TaskKeys.PARAM_DATA] - del task[TaskKeys.PARAM_GRAD] - task.clear() - - if hyperparams is not None: - optimizer.param_groups[param_group_id].update(hyperparams) - - grad_numel = param_grad.numel() - if grad_numel > max_grad_numel: - error_msg = ( - f"Gradient size {grad_numel} exceeds pre-allocated buffer size {max_grad_numel}. " - f"This indicates insufficient buffer allocation. Please increase max_grad_numel parameter.") - result_queue.put({"error": error_msg}) - break - - param_grad_cpu = grad_buffer[:grad_numel].view_as(param_grad) - if grad_scale == 1.0: - param_grad_cpu.copy_(param_grad, non_blocking=False) - else: - # A step submitted during backward gets the gradient before the loss scale is removed. - torch.mul(param_grad, grad_scale, out=param_grad_cpu) - - fp32_param = torch.nn.Parameter(param_data) - fp32_param.grad = param_grad_cpu - - optimizer.param_groups[param_group_id]['params'] = [fp32_param] - - if rollback: - logger.debug(f"Rolling back optimizer state for sub_group_id: {sub_group_id}") - optimizer.rollback_subgroup(sub_group_id) - else: - optimizer.step_subgroup(sub_group_id) - - # Send result back to main process - event_type = EventTypes.ROLLBACK if rollback else EventTypes.ADAM_STEP - result_queue.put({ - TaskKeys.PARAM_GROUP_ID: param_group_id, - TaskKeys.SUB_GROUP_ID: sub_group_id, - ResultKeys.UPDATED_PARAM: fp32_param.data, - ResultKeys.EVENT_TYPE: event_type, - }) - - # Clean up references to free memory - optimizer.param_groups[param_group_id]['params'] = [] - del param_grad_cpu, fp32_param.grad, fp32_param, param_grad, param_data - - except KeyError as e: - error_msg = f"Missing required task key: {e}" - logger.error(error_msg) - result_queue.put({"error": error_msg}) - break - except Exception as e: - error_msg = f"Unexpected error in worker process: {e}" - logger.error(error_msg) - result_queue.put({"error": error_msg}) - break - # Clean up the reusable host buffer. unpin_memory is a no-op on torch and - # frees immediately on the native backend. - if 'grad_buffer' in locals(): - if pin_memory: - get_accelerator().unpin_memory(grad_buffer) - del grad_buffer - logger.debug("Cleaned up worker grad buffer") +def _kernel_args(hyperparams: Dict[str, Any]): + beta1, beta2 = hyperparams["betas"] + return (hyperparams["lr"], beta1, beta2, hyperparams["eps"], hyperparams["weight_decay"], + hyperparams["bias_correction"]) - logger.debug("Worker process terminated") +class SuperOffloadCPUAdam: + """DeepSpeed's CPU Adam for the CPU sub-groups, run by one thread of the training process on the cores given to + CPU Adam. -class SuperOffloadCPUOptimizer: + The native kernels release the GIL, so the training thread goes on queueing backward kernels while a sub-group + steps. The thread runs what is submitted in order. The caller keeps the Adam state of each sub-group and passes it + to each step. + """ - def __init__(self, - optimizer_config: Dict[str, Any], - cpuadam_cores_perc: float = 0.8, - max_grad_numel: int = 1000000, - pin_memory: bool = True, - adamw_mode: bool = True, - fp32_optimizer_states: bool = True, - grad_dtype: torch.dtype = torch.float32) -> None: + def __init__(self, defaults: Dict[str, Any], adamw_mode: bool, fp32_optimizer_states: bool, + cpuadam_cores_perc: float) -> None: + """defaults: lr, betas, eps and weight_decay to create the native optimizer with; each step passes its own.""" if not 0 < cpuadam_cores_perc <= 1: raise ValueError("cpuadam_cores_perc must be between 0 and 1") - - self.max_grad_numel = max_grad_numel - self.mp_context = mp.get_context('spawn') - self.param_queue = self.mp_context.SimpleQueue() - self.result_queue = self.mp_context.SimpleQueue() - - self.cpuadam_process = self.mp_context.Process( - target=superoffload_optimizer_worker, - args=(self.param_queue, self.result_queue, optimizer_config, max_grad_numel, pin_memory, adamw_mode, - fp32_optimizer_states, grad_dtype), - daemon=True, - ) - self.cpuadam_process.start() - - # Set CPU affinity for better performance isolation - self._set_cpu_affinity(cpuadam_cores_perc) - - def _set_cpu_affinity(self, cpuadam_cores_perc: float) -> None: - """ - Set CPU affinity for the main (Pytorch) process and worker (CPU Adam) process. - - Args: - cpuadam_cores_perc: Percentage of cores to allocate to the worker (CPU Adam) process - """ try: - current_process = psutil.Process() - all_cores = current_process.cpu_affinity() - num_cores = len(all_cores) - - split_idx = int((1 - cpuadam_cores_perc) * num_cores) - pt_cores = all_cores[:split_idx] - cpuadam_cores = all_cores[split_idx:] - - # Set affinity for main process (PyTorch) - current_process.cpu_affinity(pt_cores) - - # Set affinity for optimizer process (CPU Adam) - optimizer_process = psutil.Process(self.cpuadam_process.pid) - optimizer_process.cpu_affinity(cpuadam_cores) - - logger.debug(f"Set CPU affinity - PyTorch cores: {pt_cores}, " - f"Optimizer cores: {cpuadam_cores}") - - except (psutil.AccessDenied, psutil.NoSuchProcess, AttributeError) as e: - logger.debug(f"Could not set CPU affinities for superoffload optimizer process: {e}") - except Exception as e: - logger.warning(f"Unexpected error setting CPU affinity: {e}") - - def async_step(self, - param_group_id: int, - sub_group_id: int, - fp32_param: torch.Tensor, - fp32_grad: torch.Tensor, - rollback: bool = False, - hyperparams: Dict[str, Any] = None, - grad_scale: float = 1.0) -> None: - """ - Queue parameter for optimization in the worker process. - - hyperparams (lr, betas, eps, weight_decay) replace those of the param group in the worker before the step. - """ - if not self.cpuadam_process.is_alive(): - raise RuntimeError("Worker process is not alive") - - task = { - TaskKeys.PARAM_DATA: fp32_param, - TaskKeys.PARAM_GRAD: fp32_grad, - TaskKeys.PARAM_GROUP_ID: param_group_id, - TaskKeys.SUB_GROUP_ID: sub_group_id, - TaskKeys.ROLLBACK: rollback, - } - if hyperparams is not None: - task[TaskKeys.HYPERPARAMS] = hyperparams + cores = sorted(os.sched_getaffinity(0)) + except AttributeError: # not Linux + cores = [] + split = int((1 - cpuadam_cores_perc) * len(cores)) + # The training thread keeps the first cores and CPU Adam gets the others, so they do not compete. + _bind_thread(cores[:split]) + self._thread = ThreadPoolExecutor(max_workers=1, + thread_name_prefix="SuperOffloadCPUAdam", + initializer=_bind_thread, + initargs=(cores[split:], )) + + self.fp32_optimizer_states = fp32_optimizer_states + self._ops = CPUAdamBuilder().load() + self._opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + self._ops.create_adam(self._opt_id, defaults["lr"], defaults["betas"][0], defaults["betas"][1], + defaults["eps"], defaults["weight_decay"], adamw_mode, should_log_le("info")) + # Holds a gradient with the loss scale removed, or one in fp32 for adam_rollback. + self._grad_buffer = None + + def submit(self, fn: Callable, *args) -> Future: + """Run fn(*args) in the CPU Adam thread, after what was submitted before it.""" + return self._thread.submit(fn, *args) + + def step(self, + state: Dict[str, Any], + weights: torch.Tensor, + grad: torch.Tensor, + hyperparams: Dict[str, Any], + grad_scale: float = 1.0) -> None: + """Step the weights on grad * grad_scale, as DeepSpeedCPUAdam does; state starts empty.""" + if not state: + state["step"] = 0 + state_dtype = torch.float if self.fp32_optimizer_states else weights.dtype + state["exp_avg"] = torch.zeros_like(weights, dtype=state_dtype) + state["exp_avg_sq"] = torch.zeros_like(weights, dtype=state_dtype) + state["step"] += 1 if grad_scale != 1.0: - task[TaskKeys.GRAD_SCALE] = grad_scale - self.param_queue.put(task) - - def _request(self, task: Dict[str, Any], event_type: str) -> Dict[str, Any]: - """Send a command to the worker process and wait for its reply. Call it with no step in flight.""" - self.param_queue.put(task) - while self.result_queue.empty(): - if not self.cpuadam_process.is_alive(): - raise RuntimeError("Worker process is not alive") - time.sleep(0.001) - return self.get_result(expected_event_type=event_type) - - def get_state(self) -> Dict[int, Dict[str, Any]]: - """Return the Adam state of each sub-group that has taken a step, keyed by sub-group id.""" - result = self._request({TaskKeys.COMMAND: EventTypes.GET_STATE}, EventTypes.GET_STATE) - return result[ResultKeys.STATE] - - def set_state(self, state: Dict[int, Dict[str, Any]]) -> None: - """Replace the Adam state of the worker, for example with the state loaded from a checkpoint.""" - self._request({TaskKeys.COMMAND: EventTypes.SET_STATE, TaskKeys.STATE: state}, EventTypes.SET_STATE) - - def get_result(self, expected_event_type: str = None) -> Optional[Dict[str, Any]]: - """ - Get result from worker process with optional event type validation. - - Args: - expected_event_type (str, optional): Expected event type ('adam_step' or 'rollback'). - If provided, validates that the result matches. - """ - if self.result_queue.empty(): - return None - - result = self.result_queue.get() - - if "error" in result: - raise RuntimeError(f"Error in worker process: {result['error']}") - - # Validate event type if expected_event_type is provided - if expected_event_type is not None: - result_event_type = result.get(ResultKeys.EVENT_TYPE) - if result_event_type != expected_event_type: - raise RuntimeError(f"Event type mismatch: expected '{expected_event_type}', got '{result_event_type}'") - - return result + grad = self._unscaled(grad, weights.dtype, grad_scale) + # adam_update reads bf16 gradients of fp32 weights as they are, so they go to the CPU at half the bytes. + self._ops.adam_update(self._opt_id, state["step"], *_kernel_args(hyperparams), weights, grad, state["exp_avg"], + state["exp_avg_sq"]) + + def rollback(self, + state: Dict[str, Any], + weights: torch.Tensor, + grad: torch.Tensor, + hyperparams: Dict[str, Any], + grad_scale: float = 1.0) -> None: + """Undo the last step of the weights, which used this gradient, grad_scale and these hyperparameters.""" + # adam_rollback reads gradients in the weights' dtype only. + self._ops.adam_rollback(self._opt_id, state["step"], *_kernel_args(hyperparams), weights, + self._unscaled(grad, weights.dtype, grad_scale), state["exp_avg"], state["exp_avg_sq"]) + state["step"] -= 1 + + def _unscaled(self, grad: torch.Tensor, dtype: torch.dtype, grad_scale: float) -> torch.Tensor: + if grad_scale == 1.0 and grad.dtype == dtype: + return grad + buffer = self._grad_buffer + if buffer is None or buffer.dtype != dtype or buffer.numel() < grad.numel(): + buffer = self._grad_buffer = torch.empty(grad.numel(), dtype=dtype) + unscaled = buffer[:grad.numel()].view_as(grad) + # A step submitted during backward gets the gradient before the loss scale is removed. + torch.mul(grad, grad_scale, out=unscaled) + return unscaled def close(self) -> None: - """ - Shutdown the worker process gracefully. - - Sends termination signal to worker and waits for clean shutdown. - If the process doesn't terminate within the timeout, it will be forcefully killed. - """ - if not self.cpuadam_process.is_alive(): - logger.debug("Worker process already terminated") - return - - # Send termination signal - self.param_queue.put(None) - - # Wait for graceful shutdown - self.cpuadam_process.join(timeout=5) - - if self.cpuadam_process.is_alive(): - logger.warning("Optimizer process did not terminate cleanly within timeout, " - "forcefully terminating") - self.cpuadam_process.terminate() - self.cpuadam_process.join(timeout=2) - - # Last resort: kill the process - if self.cpuadam_process.is_alive(): - logger.error("Failed to terminate optimizer process, killing it") - self.cpuadam_process.kill() - self.cpuadam_process.join() - - logger.debug("SuperOffload CPU optimizer closed successfully") + """Finish what was submitted, stop the thread and free the native optimizer.""" + self._thread.shutdown() + self._ops.destroy_adam(self._opt_id) diff --git a/tests/unit/v1/pin_memory/test_offload_route.py b/tests/unit/v1/pin_memory/test_offload_route.py index ce50bff43a97..64925b17e5bc 100644 --- a/tests/unit/v1/pin_memory/test_offload_route.py +++ b/tests/unit/v1/pin_memory/test_offload_route.py @@ -132,20 +132,3 @@ def _recording_synchronize(): gc.collect() # Released once the copies are known to have completed. assert begin not in manager._ranges - - -def test_superoffload_grad_buffer_unpinned_when_disabled(): - """offload_optimizer.pin_memory=False must allocate a regular CPU buffer.""" - from deepspeed.runtime.superoffload.superoffload_utils import _allocate_worker_grad_buffer - buffer = _allocate_worker_grad_buffer(32, pin_memory=False) - assert get_accelerator().is_pinned(buffer) is False - - -def test_superoffload_grad_buffer_pinned_when_enabled(monkeypatch): - """offload_optimizer.pin_memory=True keeps the current pinned allocation.""" - _require_native(monkeypatch) - from deepspeed.runtime.superoffload.superoffload_utils import _allocate_worker_grad_buffer - accel = get_accelerator() - buffer = _allocate_worker_grad_buffer(32, pin_memory=True) - assert accel.is_pinned(buffer) is True - assert accel.unpin_memory(buffer) is True diff --git a/tests/unit/v1/zero/test_invalid_grad_norm.py b/tests/unit/v1/zero/test_invalid_grad_norm.py index 152161fbc736..4315a45ce8f8 100644 --- a/tests/unit/v1/zero/test_invalid_grad_norm.py +++ b/tests/unit/v1/zero/test_invalid_grad_norm.py @@ -154,16 +154,17 @@ def test_bf16_optimizer_reports_invalid_step_and_recovers(monkeypatch): def test_superoffload_rolls_back_only_subgroups_submitted_this_step(): optimizer = object.__new__(SuperOffloadOptimizer_Stage3) - optimizer._submitted_cpu_sub_groups = {1: ({"lr": 0.1}, 1.0)} + optimizer._submitted_cpu_sub_groups = {1: SimpleNamespace(hyperparams={"lr": 0.1}, grad_scale=1.0)} + optimizer._half_cpu_grads = {} optimizer.sub_group_to_group_id = {0: 10, 1: 11} parameter = SimpleNamespace(data=torch.tensor([1.0]), grad=SimpleNamespace(data=torch.tensor([2.0]))) optimizer.fp32_partitioned_groups_flat = [parameter, parameter] rollbacks = [] - optimizer._sync_cpu_optimizer_step = lambda *args, **kwargs: rollbacks.append((args, kwargs)) + optimizer._cpu_step_and_wait = lambda *args, **kwargs: rollbacks.append((args, kwargs)) optimizer._rollback_submitted_cpu_sub_groups() - assert [args[1] for args, _ in rollbacks] == [1] + assert [args[0] for args, _ in rollbacks] == [1] assert rollbacks[0][1]["rollback"] assert optimizer._submitted_cpu_sub_groups == {} diff --git a/tests/unit/v1/zero/test_superoffload.py b/tests/unit/v1/zero/test_superoffload.py index bbcb80970254..b0248bb4ca34 100644 --- a/tests/unit/v1/zero/test_superoffload.py +++ b/tests/unit/v1/zero/test_superoffload.py @@ -2,8 +2,11 @@ # DeepSpeed Team """SuperOffload against ZeRO-3 CPU offload, which applies the same DeepSpeedCPUAdam update in the main process.""" +import time + import pytest import torch +import torch.utils.checkpoint import deepspeed import deepspeed.comm as dist @@ -20,7 +23,8 @@ class PartlyUsedModel(torch.nn.Module): - """`a` is always used. `b` and `c` are used, skipped, or multiplied by 0, which gives them a zero gradient.""" + """`a` is always used. `b` and `c` are used, skipped, multiplied by 0, which gives them a zero gradient, or used + twice, which gives them two gradients in one step.""" def __init__(self): super().__init__() @@ -36,10 +40,25 @@ def forward(self, x, uses): h = h + getattr(self, name)(x) elif use == "zero": h = h + 0 * getattr(self, name)(x) + elif use in ("twice", "twice_in_one_backward"): + layer = getattr(self, name) + if use == "twice": + # The backward of a reentrant checkpoint runs on its own, so the param's gradient hook fires twice. + again = torch.utils.checkpoint.checkpoint(layer, h, use_reentrant=True) + else: + again = layer(h) + h = h + layer(x) + again return h.float().pow(2).mean() -def _config(super_offload, clip=0.0, ga=1, precision="bf16", ratio=1.0, sub_group_size=1e9, **optimizer_params): +def _config(super_offload, + clip=0.0, + ga=1, + precision="bf16", + ratio=1.0, + sub_group_size=1e9, + reduce_bucket_size=None, + **optimizer_params): optimizer_params.setdefault("lr", 3e-3) optimizer_params.setdefault("weight_decay", 0.01) optimizer_type = optimizer_params.pop("type", "AdamW") @@ -62,6 +81,8 @@ def _config(super_offload, clip=0.0, ga=1, precision="bf16", ratio=1.0, sub_grou }, }, } + if reduce_bucket_size is not None: + config["zero_optimization"]["reduce_bucket_size"] = reduce_bucket_size if precision == "fp16": config["fp16"] = {"enabled": True, "initial_scale_power": 24} elif precision != "fp32": @@ -93,8 +114,8 @@ def _masters(engine): def _close(engine): - if hasattr(engine.optimizer, "superoffload_cpu_optimizer"): - engine.optimizer.superoffload_cpu_optimizer.close() + if hasattr(engine.optimizer, "cpu_adam"): + engine.optimizer.cpu_adam.close() def _train(engine, steps, uses_for, first_step=0, before_step=None): @@ -130,6 +151,13 @@ def _b_unused_on_even_steps(unused): return lambda step, micro_step: {"b": "use" if step % 2 else unused} +def _b_used_twice(unused): + # ZeRO-3 keeps only the last gradient when the hook of a param fires twice in a step, so the reference uses the + # param twice in one backward, where autograd sums the two gradients. + use = "twice_in_one_backward" if unused == "zero" else "twice" + return lambda step, micro_step: {"b": use} + + def _b_and_c_in_one_micro_step(unused): return lambda step, micro_step: {"b": unused, "c": "use"} if micro_step else {"b": "use", "c": unused} @@ -141,6 +169,10 @@ def _b_and_c_in_one_micro_step(unused): # Clipping is on but never triggers, so the CPU steps still run during backward and must stay exact. "unused_param_clip": (dict(clip=100.0), _b_unused_on_even_steps, True), "grad_accumulation": (dict(ga=2), _b_and_c_in_one_micro_step, True), + # The second gradient adds to the first in the gradient accumulation dtype, bf16 or fp32. ZeRO-3 takes a param + # once per reduce bucket, so a small bucket keeps the two gradients apart. + "param_used_twice": (dict(reduce_bucket_size=1000), _b_used_twice, True), + "param_used_twice_fp32": (dict(reduce_bucket_size=1000, precision="fp32"), _b_used_twice, True), "adam": (dict(type="Adam", adam_w_mode=False), lambda unused: _all_used, True), "no_bias_correction": (dict(bias_correction=False), lambda unused: _all_used, True), "bf16_master_weights": (dict(precision="bf16_master"), lambda unused: _all_used, True), @@ -160,8 +192,6 @@ def _b_and_c_in_one_micro_step(unused): class TestSuperOffloadMatchesZeRO3(DistributedTest): world_size = 1 - # SuperOffload starts its CPU worker process, which a daemonic test process cannot do. - non_daemonic_procs = True @pytest.mark.parametrize("case", list(CASES)) def test(self, case): @@ -197,7 +227,6 @@ def _change_hyperparams(engine, step): class TestSuperOffloadFollowsZeRO3(DistributedTest): world_size = 1 - non_daemonic_procs = True def test_hyperparameter_changes(self): """Momentum and weight decay schedules change betas and weight_decay of the param groups between steps.""" @@ -224,21 +253,42 @@ def test_fp16_loss_scale_increase(self): def test_cpu_steps_run_during_backward(self): engine = _engine(_config(super_offload=True)) - cpu_optimizer = engine.optimizer.superoffload_cpu_optimizer - submitted = [] - async_step = cpu_optimizer.async_step + cpu_adam = engine.optimizer.cpu_adam + calls = [] - def recording_async_step(*args, **kwargs): - submitted.append(kwargs.get("rollback", False)) - async_step(*args, **kwargs) + def recording(name): + update = getattr(cpu_adam, name) - cpu_optimizer.async_step = recording_async_step + def record(state, weights, grad, *args): + calls.append((name, grad.dtype)) + update(state, weights, grad, *args) + + return record + + cpu_adam.step, cpu_adam.rollback = recording("step"), recording("rollback") x = torch.randn(4, HIDDEN).to(engine.device, torch.bfloat16) engine.backward(engine(x, {})) cpu_sub_groups = list(engine.optimizer.subgroup_to_device.values()).count("cpu") - assert submitted == [False] * cpu_sub_groups + # The sub-groups step once their gradients reach the CPU, without waiting for step(), on the bf16 gradients + # as they went to the CPU. + deadline = time.monotonic() + 30 + while len(calls) < cpu_sub_groups and time.monotonic() < deadline: + time.sleep(0.01) + assert calls == [("step", torch.bfloat16)] * cpu_sub_groups engine.step() - assert submitted == [False] * cpu_sub_groups + assert calls == [("step", torch.bfloat16)] * cpu_sub_groups + # CPU Adam steps the weights where they are, without moving them to shared memory. + assert not any(t.is_shared() for t in engine.optimizer.fp32_partitioned_groups_flat) + _close(engine) + + def test_cpu_buffers_pinned(self): + config = _config(super_offload=True) + config["zero_optimization"]["offload_optimizer"]["pin_memory"] = True + engine = _engine(config) + # The gradient copies of backward and the weight copies of step() stay asynchronous only from page-locked + # memory; ZeRO-3 pins the gradients but not the fp32 weights. + assert all(t.is_pinned() and t.grad.is_pinned() for t in engine.optimizer.fp32_partitioned_groups_flat) + assert engine.optimizer._lp_staging_stream is not None _close(engine) @@ -258,7 +308,6 @@ def forward(self, tokens, inject): class TestSuperOffloadSkippedStep(DistributedTest): world_size = 1 - non_daemonic_procs = True def test_sub_group_with_finite_grads(self): """A step skipped for an inf gradient elsewhere undoes the CPU sub-group's step; its Adam state must stay @@ -286,7 +335,6 @@ def test_sub_group_with_finite_grads(self): class TestSuperOffloadTwoRanks(DistributedTest): world_size = 2 - non_daemonic_procs = True def test_unused_param(self): engine = _engine(_config(super_offload=False)) @@ -316,10 +364,33 @@ def _resume_matches_uninterrupted(config, checkpoint_dir, two_groups): class TestSuperOffloadCheckpoint(DistributedTest): world_size = 2 - non_daemonic_procs = True # Betas of 0.5 make beta**t exact, so the resume must be bit-exact (CPU Adam recomputes beta**t on resume). @pytest.mark.parametrize("two_groups", [False, True]) def test_resume(self, tmpdir, two_groups): config = _config(super_offload=True, sub_group_size=4500, betas=[0.5, 0.5]) _resume_matches_uninterrupted(config, str(tmpdir), two_groups) + + +def test_cpu_adam_steps_and_rolls_back_on_bf16_grads(): + """A step on a bf16 gradient matches DeepSpeedCPUAdam on the same gradient in fp32, and the rollback undoes it.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + from deepspeed.runtime.superoffload.superoffload_utils import SuperOffloadCPUAdam + + hyperparams = dict(lr=1e-2, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, bias_correction=True) + torch.manual_seed(0) + weights, grad = torch.randn(4099), torch.randn(4099).to(torch.bfloat16) + expected = torch.nn.Parameter(weights.clone()) + expected.grad = grad.float() + DeepSpeedCPUAdam([expected], **hyperparams).step() + original = weights.clone() + # All the cores go to CPU Adam, which leaves the affinity of this thread as it is. + cpu_adam = SuperOffloadCPUAdam(hyperparams, adamw_mode=True, fp32_optimizer_states=True, cpuadam_cores_perc=1.0) + try: + state = {} + cpu_adam.submit(cpu_adam.step, state, weights, grad, hyperparams).result() + assert torch.equal(weights, expected.data) + cpu_adam.submit(cpu_adam.rollback, state, weights, grad, hyperparams).result() + torch.testing.assert_close(weights, original, rtol=0, atol=1e-6) + finally: + cpu_adam.close() diff --git a/tests/unit/v1/zero/test_zero_partial_offload.py b/tests/unit/v1/zero/test_zero_partial_offload.py index 9b2c98cf1d79..dd0c8064d0d1 100644 --- a/tests/unit/v1/zero/test_zero_partial_offload.py +++ b/tests/unit/v1/zero/test_zero_partial_offload.py @@ -87,8 +87,8 @@ def _train(engine, steps, first_step=0): def _close(engine): - if hasattr(engine.optimizer, "superoffload_cpu_optimizer"): - engine.optimizer.superoffload_cpu_optimizer.close() + if hasattr(engine.optimizer, "cpu_adam"): + engine.optimizer.cpu_adam.close() def _device_params(engine): @@ -106,8 +106,6 @@ def _adam_state(engine): class TestPartialOffloadHyperparameters(DistributedTest): world_size = 1 - # SuperOffload starts its CPU worker process, which a daemonic test process cannot do. - non_daemonic_procs = True @pytest.mark.parametrize("super_offload", [False, True]) def test_first_step_at_zero_lr(self, super_offload): From 2ce5bcc7c5b0a1b28ad9d55820028932a11b55e1 Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Fri, 6 Nov 2026 07:59:22 +0000 Subject: [PATCH 08/10] SuperOffload: copy large gradients to the CPU in pieces on two streams 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 --- .../superoffload/superoffload_stage3.py | 47 ++++++++++++++++--- tests/unit/v1/zero/test_superoffload.py | 28 ++++++++--- 2 files changed, 62 insertions(+), 13 deletions(-) diff --git a/deepspeed/runtime/superoffload/superoffload_stage3.py b/deepspeed/runtime/superoffload/superoffload_stage3.py index 48666660982b..8b2bc1c3d084 100644 --- a/deepspeed/runtime/superoffload/superoffload_stage3.py +++ b/deepspeed/runtime/superoffload/superoffload_stage3.py @@ -19,6 +19,9 @@ from deepspeed.accelerator import get_accelerator OPTIMIZER_STEP_TIMER = 'optimizer_step' +# Larger gradients go to the CPU in pieces of this size (see _copy_grad_to_cpu). On B200 over PCIe Gen5, 16 MiB pieces +# copy at 98% of the rate of whole gradients; smaller ones lose more to the switches between pieces. +_GRAD_COPY_PIECE_BYTES = 16 << 20 def _validate_superoffload_accelerator(): @@ -105,6 +108,10 @@ def __init__( self.fp32_partitioned_groups_flat[sub_group_id].numel(), dtype=torch.bfloat16) self._pin_cpu_buffers() self._grad_copy_stream = get_accelerator().Stream() + # The pieces of a large gradient take turns on the copy stream and a second one, each with an event + # (_copy_grad_to_cpu). + self._grad_copy_streams = (self._grad_copy_stream, get_accelerator().Stream()) + self._grad_copy_events = (get_accelerator().Event(), get_accelerator().Event()) if self.speculative_cpu_step and self._cpu_sub_group_weights_pinned(): # Default priority is the lowest, so the conversions never delay the backward kernels. self._lp_staging_stream = get_accelerator().Stream() @@ -341,12 +348,7 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L cpu_grad_tensor = self._worker_grad(i).narrow(0, dest_offset, grad_buffer.numel()) grad_to_copy = grad_buffer.to(dtype=cpu_grad_tensor.dtype) - copy_stream = self._grad_copy_stream - copy_stream.wait_stream(get_accelerator().current_stream()) - with get_accelerator().stream(copy_stream): - cpu_grad_tensor.copy_(grad_to_copy, non_blocking=True) - # The copy reads the gradient later, so keep the allocator from handing its memory out before then. - grad_to_copy.record_stream(copy_stream) + self._copy_grad_to_cpu(cpu_grad_tensor, grad_to_copy) self.sub_group_grads_ready[i].add(param.ds_id) if len(self.sub_group_grads_ready[i]) == self.sub_group_to_param_num[i]: @@ -395,6 +397,39 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L param.grad.record_stream(get_accelerator().current_stream()) param.grad = None + def _copy_grad_to_cpu(self, cpu_grad_tensor, grad_to_copy): + """Queue the copy of a gradient to the CPU on the gradient copy stream, behind the kernels that computed it. + + The copy engine keeps serving a stream while it has a copy ready, so a copy to the CPU on another stream, such + as the one .item() makes in backward, waits behind all the gradient copies queued before it: up to 28 ms with + OPT-30B. A large gradient goes in pieces that take turns on two streams, each waiting for the one before, so + that neither stream has a copy ready while a piece runs and the copies of other streams go in between. + """ + copy_stream = self._grad_copy_stream + copy_stream.wait_stream(get_accelerator().current_stream()) + piece = _GRAD_COPY_PIECE_BYTES // grad_to_copy.element_size() + if grad_to_copy.numel() <= piece: + with get_accelerator().stream(copy_stream): + cpu_grad_tensor.copy_(grad_to_copy, non_blocking=True) + # The copy reads the gradient later, so keep the allocator from handing its memory out before then. + grad_to_copy.record_stream(copy_stream) + return + streams = self._grad_copy_streams + events = self._grad_copy_events + previous = None + for turn, start in enumerate(range(0, grad_to_copy.numel(), piece)): + stream = streams[turn % 2] + if previous is not None: + stream.wait_event(previous) + with get_accelerator().stream(stream): + cpu_grad_tensor[start:start + piece].copy_(grad_to_copy[start:start + piece], non_blocking=True) + previous = events[turn % 2] + previous.record(stream) + # What waits for the copy stream, as step() and the CPU steps do, then waits for the last piece too. + copy_stream.wait_event(previous) + for stream in streams: + grad_to_copy.record_stream(stream) + @instrument_w_nvtx def step(self, closure=None): """ diff --git a/tests/unit/v1/zero/test_superoffload.py b/tests/unit/v1/zero/test_superoffload.py index b0248bb4ca34..0a7f0cd67bd1 100644 --- a/tests/unit/v1/zero/test_superoffload.py +++ b/tests/unit/v1/zero/test_superoffload.py @@ -26,11 +26,12 @@ class PartlyUsedModel(torch.nn.Module): """`a` is always used. `b` and `c` are used, skipped, multiplied by 0, which gives them a zero gradient, or used twice, which gives them two gradients in one step.""" - def __init__(self): + def __init__(self, hidden=HIDDEN): super().__init__() - self.a = torch.nn.Linear(HIDDEN, HIDDEN) - self.b = torch.nn.Linear(HIDDEN, HIDDEN) - self.c = torch.nn.Linear(HIDDEN, HIDDEN) + self.hidden = hidden + self.a = torch.nn.Linear(hidden, hidden) + self.b = torch.nn.Linear(hidden, hidden) + self.c = torch.nn.Linear(hidden, hidden) def forward(self, x, uses): h = torch.nn.functional.gelu(self.a(x)) @@ -94,9 +95,9 @@ def _config(super_offload, return config -def _engine(config, two_groups=False): +def _engine(config, two_groups=False, hidden=HIDDEN): torch.manual_seed(0) - model = PartlyUsedModel() + model = PartlyUsedModel(hidden) params = model.parameters() if two_groups: # Weights with weight decay and biases without, as most training scripts set them up. @@ -135,7 +136,7 @@ def _train(engine, steps, uses_for, first_step=0, before_step=None): if before_step is not None: before_step(engine, step) for micro_step in range(engine.gradient_accumulation_steps()): - x = torch.randn(4, HIDDEN, generator=generator).to(engine.device, dtype) + x = torch.randn(4, engine.module.hidden, generator=generator).to(engine.device, dtype) loss = engine(x, uses_for(step, micro_step)) engine.backward(loss) engine.step() @@ -237,6 +238,19 @@ def test_hyperparameter_changes(self): _close(engine) _assert_same(results[1], results[0], exact=True) + def test_large_grads(self): + """The bf16 gradient of a 4090 x 4090 weight, 31.9 MiB, goes to the CPU in two pieces of at most 16 MiB. A CPU + step that did not wait for the second piece would read part of the previous gradient. Only to pinned memory + are the copies asynchronous.""" + results = [] + for super_offload in (False, True): + config = _config(super_offload=super_offload) + config["zero_optimization"]["offload_optimizer"]["pin_memory"] = True + engine = _engine(config, hidden=4090) + results.append(_train(engine, 3, _all_used)) + _close(engine) + _assert_same(results[1], results[0], exact=True) + def test_fp16_loss_scale_increase(self): """Dynamic loss scaling raises the scale in step(), after the CPU steps taken during backward used the old one; those steps are undone and redone, exact up to rounding, which fp16 weights amplify (1.4e-5 after 6 From ad833ac8922a7c733102177d3d0969e28e4afa51 Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Sun, 8 Nov 2026 02:24:29 +0000 Subject: [PATCH 09/10] CPU Adam: scale gradients as the step reads them, change a step to another 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 --- csrc/adam/cpu_adam.cpp | 24 ++++- csrc/adam/cpu_adam_impl.cpp | 103 +++++++++++++++++++- csrc/includes/cpu_adam.h | 139 ++++++++++++++++++++++++++- csrc/includes/simd.h | 13 +++ tests/unit/ops/adam/test_cpu_adam.py | 90 +++++++++++++++++ 5 files changed, 362 insertions(+), 7 deletions(-) diff --git a/csrc/adam/cpu_adam.cpp b/csrc/adam/cpu_adam.cpp index 90c7c188075b..1b899bf14f4c 100644 --- a/csrc/adam/cpu_adam.cpp +++ b/csrc/adam/cpu_adam.cpp @@ -8,11 +8,29 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { using namespace pybind11::literals; - // The update and the rollback release the GIL, so that a CPU Adam step in another thread lets - // the training thread run. + // The update, the rescale and the rollback release the GIL, so that a CPU Adam step in another + // thread lets the training thread run. m.def("adam_update", &ds_adam_step, - "DeepSpeed CPU Adam update (C++)", + "DeepSpeed CPU Adam update (C++); grad_scale multiplies each gradient as it is read", + "optimizer_id"_a, + "step"_a, + "lr"_a, + "beta1"_a, + "beta2"_a, + "epsilon"_a, + "weight_decay"_a, + "bias_correction"_a, + "params"_a, + "grads"_a, + "exp_avg"_a, + "exp_avg_sq"_a, + "grad_scale"_a = 1.0f, + pybind11::call_guard()); + m.def("adam_rescale_step", + &ds_adam_rescale_step, + "DeepSpeed CPU Adam: turn the last step, taken on grads * old_scale, into the step on " + "grads * new_scale (C++)", pybind11::call_guard()); m.def("adam_rollback", &ds_adam_rollback, diff --git a/csrc/adam/cpu_adam_impl.cpp b/csrc/adam/cpu_adam_impl.cpp index 0dc5932d7c75..ad13629cffbc 100644 --- a/csrc/adam/cpu_adam_impl.cpp +++ b/csrc/adam/cpu_adam_impl.cpp @@ -64,6 +64,7 @@ void Adam_Optimizer::Step_1(ds_params_precision_t* _params, #pragma omp parallel for if (parallel) for (size_t k = t; k < offset; k++) { float grad = (float)grads[k]; + if (_grad_scale != 1.0f) { grad = grad * _grad_scale; } float param = (float)_params[k]; float momentum = _exp_avg[k]; float variance = _exp_avg_sq[k]; @@ -274,7 +275,8 @@ int ds_adam_step(int optimizer_id, torch::Tensor& params, torch::Tensor& grads, torch::Tensor& exp_avg, - torch::Tensor& exp_avg_sq) + torch::Tensor& exp_avg_sq, + float grad_scale) { auto params_c = params.contiguous(); auto grads_c = grads.contiguous(); @@ -284,13 +286,110 @@ int ds_adam_step(int optimizer_id, std::shared_ptr opt = std::static_pointer_cast(s_optimizers[optimizer_id]); opt->IncrementStep(step, beta1, beta2); - opt->update_state(lr, epsilon, weight_decay, bias_correction); + opt->update_state(lr, epsilon, weight_decay, bias_correction, grad_scale); invoke(opt, params_c, grads_c, exp_avg_c, exp_avg_sq_c, params_c.numel()); return 0; } +template +void Adam_Optimizer::Rescale_Step(float* _params, + ds_grad_precision_t* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t _param_size, + float old_scale, + float new_scale) +{ + if (_weight_decay > 0 && !_adamw_mode) { + // Adam's L2 decay adds the params of before the step to the gradient, which are gone. + throw std::runtime_error( + "Adam rescale step supports AdamW only, or Adam without weight decay"); + } + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__) + Rescale_Step_AVX<8>( + &rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size, old_scale, new_scale); +#endif + if (_param_size <= rounded_size) return; + float momentum_shift = (1 - _betta1) * (new_scale - old_scale); + float variance_shift = (1 - _betta2) * (new_scale * new_scale - old_scale * old_scale); + float step_size = -1 * _alpha / _bias_correction1; +#pragma omp parallel for + for (size_t k = rounded_size; k < _param_size; k++) { + float grad = (float)grads[k]; + float momentum = _exp_avg[k]; + float variance = _exp_avg_sq[k]; + float param = _params[k]; + param = momentum / (sqrt(variance) * _bias_correction2 + _eps) * -step_size + param; + momentum = grad * momentum_shift + momentum; + // std::max(x, 0) keeps a NaN x. + variance = std::max(grad * grad * variance_shift + variance, 0.0f); + param = momentum / (sqrt(variance) * _bias_correction2 + _eps) * step_size + param; + _params[k] = param; + _exp_avg[k] = momentum; + _exp_avg_sq[k] = variance; + } +} + +int ds_adam_rescale_step(int optimizer_id, + size_t step, + float lr, + float beta1, + float beta2, + float epsilon, + float weight_decay, + bool bias_correction, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg, + torch::Tensor& exp_avg_sq, + float old_scale, + float new_scale) +{ + TORCH_CHECK(params.scalar_type() == torch::kFloat32 && + exp_avg.scalar_type() == torch::kFloat32 && + exp_avg_sq.scalar_type() == torch::kFloat32, + "Adam rescale step supports float32 params and states only"); + TORCH_CHECK(grads.scalar_type() == torch::kFloat32 || grads.scalar_type() == torch::kBFloat16, + "Adam rescale step supports float32 and bfloat16 gradients only"); + TORCH_CHECK(params.is_contiguous() && grads.is_contiguous() && exp_avg.is_contiguous() && + exp_avg_sq.is_contiguous(), + "Adam rescale step updates its tensors in place, so they must be contiguous"); + TORCH_CHECK(grads.numel() == params.numel() && exp_avg.numel() == params.numel() && + exp_avg_sq.numel() == params.numel(), + "Adam rescale step needs as many gradients and states as params"); + + std::shared_ptr opt = + std::static_pointer_cast(s_optimizers[optimizer_id]); + // As the step being rescaled had this step and these betas, beta1^step and beta2^step stay. + opt->IncrementStep(step, beta1, beta2); + opt->update_state(lr, epsilon, weight_decay, bias_correction); + + float* params_ptr = params.data_ptr(); + float* exp_avg_ptr = exp_avg.data_ptr(); + float* exp_avg_sq_ptr = exp_avg_sq.data_ptr(); + if (grads.scalar_type() == torch::kBFloat16) { + opt->Rescale_Step(params_ptr, + (c10::BFloat16*)grads.data_ptr(), + exp_avg_ptr, + exp_avg_sq_ptr, + params.numel(), + old_scale, + new_scale); + } else { + opt->Rescale_Step(params_ptr, + grads.data_ptr(), + exp_avg_ptr, + exp_avg_sq_ptr, + params.numel(), + old_scale, + new_scale); + } + return 0; +} + void adamw_rollback_inplace(float* params, const float* grads, float* momentum, diff --git a/csrc/includes/cpu_adam.h b/csrc/includes/cpu_adam.h index 59d4a246d967..c961152a4e22 100644 --- a/csrc/includes/cpu_adam.h +++ b/csrc/includes/cpu_adam.h @@ -43,6 +43,7 @@ class Adam_Optimizer { _betta1_t(1.0), _betta2_t(1.0), _step(0), + _grad_scale(1.0f), _adamw_mode(adamw_mode) { } @@ -90,6 +91,25 @@ class Adam_Optimizer { STEP(1) STEP(4) STEP(8) +#if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__) + template + void Rescale_Step_AVX(size_t* rounded_size, + float* _params, + ds_grad_precision_t* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t param_size, + float old_scale, + float new_scale); +#endif + template + void Rescale_Step(float* _params, + ds_grad_precision_t* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t _param_size, + float old_scale, + float new_scale); inline void IncrementStep(size_t step, float beta1, float beta2) { if (beta1 != _betta1 || beta2 != _betta2) { @@ -113,11 +133,17 @@ class Adam_Optimizer { } } } - inline void update_state(float lr, float epsilon, float weight_decay, bool bias_correction) + // grad_scale multiplies the gradients as the step reads them, as unscaling and clipping does. + inline void update_state(float lr, + float epsilon, + float weight_decay, + bool bias_correction, + float grad_scale = 1.0f) { _alpha = lr; _eps = epsilon; _weight_decay = weight_decay; + _grad_scale = grad_scale; _bias_correction1 = 1.0f; _bias_correction2 = 1.0f; @@ -141,6 +167,8 @@ class Adam_Optimizer { float _bias_correction1; float _bias_correction2; + float _grad_scale; + bool _adamw_mode; }; @@ -192,6 +220,8 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, AVX_Data weight_decay4; if (_weight_decay > 0) weight_decay4.data = (_adamw_mode ? SIMD_SET(w_decay) : SIMD_SET(_weight_decay)); + AVX_Data grad_scale_4; + grad_scale_4.data = SIMD_SET(_grad_scale); new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); for (size_t t = 0; t < new_rounded_size; t += TILE) { size_t copy_size = TILE; @@ -201,6 +231,7 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { AVX_Data grad_4[span]; simd_load(grad_4, grads + i); + if (_grad_scale != 1.0f) { simd_mul(grad_4, grad_4, grad_scale_4); } AVX_Data momentum_4[span]; simd_load(momentum_4, _exp_avg + i); @@ -237,6 +268,90 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, } *rounded_size = new_rounded_size; } + +// Turns the last step of fp32 params and states, taken on grads * old_scale, into the step on +// grads * new_scale, in one pass where undoing it and stepping again take two. The moments move by +// the difference of the two gradients, and the params by the difference of the two Adam updates: +// AdamW's weight decay does not depend on the gradient. +template +void Adam_Optimizer::Rescale_Step_AVX(size_t* rounded_size, + float* _params, + ds_grad_precision_t* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t _param_size, + float old_scale, + float new_scale) +{ +#if !defined(__AVX512__) && !defined(__NEON__) + if (std::is_same_v) { return; } +#endif + AVX_Data momentum_shift_4; + momentum_shift_4.data = SIMD_SET((1 - _betta1) * (new_scale - old_scale)); + AVX_Data variance_shift_4; + variance_shift_4.data = + SIMD_SET((1 - _betta2) * (new_scale * new_scale - old_scale * old_scale)); + + AVX_Data bias2_sqrt; + bias2_sqrt.data = SIMD_SET(_bias_correction2); + + AVX_Data eps_4; + eps_4.data = SIMD_SET(_eps); + + float step_size = -1 * _alpha / _bias_correction1; + AVX_Data step_size_4; + step_size_4.data = SIMD_SET(step_size); + AVX_Data minus_step_size_4; + minus_step_size_4.data = SIMD_SET(-step_size); + + AVX_Data zero_4; + zero_4.data = SIMD_SET(0.0f); + + size_t new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); + for (size_t t = 0; t < new_rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { + AVX_Data grad_4[span]; + simd_load(grad_4, grads + i); + + AVX_Data momentum_4[span]; + simd_load(momentum_4, _exp_avg + i); + + AVX_Data variance_4[span]; + simd_load(variance_4, _exp_avg_sq + i); + + AVX_Data param_4[span]; + simd_load(param_4, _params + i); + + // The update of the last step, as Step_AVX computed it from these moments. + AVX_Data update_4[span]; + simd_sqrt(update_4, variance_4); + simd_fma(update_4, update_4, bias2_sqrt, eps_4); + simd_div(update_4, momentum_4, update_4); + simd_fma(param_4, update_4, minus_step_size_4, param_4); + + simd_fma(momentum_4, grad_4, momentum_shift_4, momentum_4); + simd_mul(grad_4, grad_4, grad_4); + simd_fma(variance_4, grad_4, variance_shift_4, variance_4); + // The shift can leave a rounding residual below 0 where the old variance was 0, + // whose sqrt is NaN. + simd_max(variance_4, zero_4, variance_4); + + simd_sqrt(update_4, variance_4); + simd_fma(update_4, update_4, bias2_sqrt, eps_4); + simd_div(update_4, momentum_4, update_4); + simd_fma(param_4, update_4, step_size_4, param_4); + + simd_store(_params + i, param_4); + simd_store(_exp_avg + i, momentum_4); + simd_store(_exp_avg_sq + i, variance_4); + } + } + *rounded_size = new_rounded_size; +} #endif #if defined(__SVE__) && defined(__ARM_FEATURE_SVE) @@ -255,6 +370,8 @@ void Adam_Optimizer::Step_SVE(size_t* rounded_size, if constexpr (std::is_same_v && std::is_same_v && std::is_same_v) { + // The SVE kernel does not scale the gradient; the scalar loop steps a scaled one. + if (_grad_scale != 1.0f) { return; } sve_adam_update(_params, grads, _exp_avg, @@ -290,6 +407,8 @@ void Adam_Optimizer::Step_RVV(size_t* rounded_size, if constexpr (std::is_same_v && std::is_same_v && std::is_same_v) { + // The RVV kernel does not scale the gradient; the scalar loop steps a scaled one. + if (_grad_scale != 1.0f) { return; } rvv_adam_update(_params, grads, _exp_avg, @@ -329,7 +448,23 @@ int ds_adam_step(int optimizer_id, torch::Tensor& params, torch::Tensor& grads, torch::Tensor& exp_avg, - torch::Tensor& exp_avg_sq); + torch::Tensor& exp_avg_sq, + float grad_scale); + +int ds_adam_rescale_step(int optimizer_id, + size_t step, + float lr, + float beta1, + float beta2, + float epsilon, + float weight_decay, + bool bias_correction, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg, + torch::Tensor& exp_avg_sq, + float old_scale, + float new_scale); int ds_adam_rollback(int optimizer_id, size_t step, diff --git a/csrc/includes/simd.h b/csrc/includes/simd.h index c9bbd63bcb88..b725dff72ac1 100644 --- a/csrc/includes/simd.h +++ b/csrc/includes/simd.h @@ -42,6 +42,8 @@ inline void writeAs(void* dst, const T& val) #define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c) #define SIMD_SQRT(x) _mm512_sqrt_ps(x) #define SIMD_DIV(x, y) _mm512_div_ps(x, y) +// Returns y when either operand is NaN. +#define SIMD_MAX(x, y) _mm512_max_ps(x, y) #define SIMD_AND(x, y) _mm512_and_ps(x, y) #define SIMD_ANDNOT(x, y) _mm512_andnot_ps(x, y) #define SIMD_OR(x, y) _mm512_or_ps(x, y) @@ -105,6 +107,8 @@ static void store_16_f32_as_bf16_nearest(__m512 v, void* data) #define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c) #define SIMD_SQRT(x) _mm256_sqrt_ps(x) #define SIMD_DIV(x, y) _mm256_div_ps(x, y) +// Returns y when either operand is NaN. +#define SIMD_MAX(x, y) _mm256_max_ps(x, y) #define SIMD_AND(x, y) _mm256_and_ps(x, y) #define SIMD_ANDNOT(x, y) _mm256_andnot_ps(x, y) #define SIMD_OR(x, y) _mm256_or_ps(x, y) @@ -128,6 +132,8 @@ static void store_16_f32_as_bf16_nearest(__m512 v, void* data) #define SIMD_FMA(x, y, c) vfmaq_f32(c, x, y) #define SIMD_SQRT(x) vsqrtq_f32(x) #define SIMD_DIV(x, y) vdivq_f32(x, y) +// Returns NaN when either operand is NaN. +#define SIMD_MAX(x, y) vmaxq_f32(x, y) #define SIMD_AND(x, y) \ vreinterpretq_f32_u32(vandq_u32(vreinterpretq_u32_f32(x), vreinterpretq_u32_f32(y))) // x86 andnot(x, y) computes ~x & y; NEON bic(a, b) computes a & ~b, so the operands swap. @@ -305,6 +311,13 @@ inline void simd_div(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) #pragma unroll for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_DIV(src_a_l[i].data, src_a_r[i].data); } } +// max(src_a_l, src_a_r) keeps a NaN of src_a_r on every ISA. +template +inline void simd_max(AVX_Data* dst, AVX_Data src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MAX(src_a_l.data, src_a_r[i].data); } +} template inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) { diff --git a/tests/unit/ops/adam/test_cpu_adam.py b/tests/unit/ops/adam/test_cpu_adam.py index 280b48c69242..422b7e1e6c3b 100644 --- a/tests/unit/ops/adam/test_cpu_adam.py +++ b/tests/unit/ops/adam/test_cpu_adam.py @@ -191,6 +191,96 @@ def test_cpu_adam_steps_fp32_params_on_bf16_grads(): op.destroy_adam(opt_id) +def test_cpu_adam_scales_grads_as_it_reads_them(): + """adam_update with grad_scale steps bit for bit as on the gradient scaled before, as unscale_and_clip_grads + scales it, for fp32 and bf16 gradients of fp32 params.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + op = CPUAdamBuilder().load() + opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + lr, beta1, beta2, eps, wd, scale = 1e-3, 0.9, 0.999, 1e-8, 0.01, 0.37 + n = 100003 # not a multiple of the SIMD block, so the scalar tail runs too + torch.manual_seed(0) + params = torch.randn(n) + for grad_dtype in (torch.bfloat16, torch.float): + grads = [torch.randn(n).to(grad_dtype) for _ in range(3)] + results = [] + for in_kernel in (True, False): + # A new native optimizer for each, as it carries beta**step over from the steps before. + op.create_adam(opt_id, lr, beta1, beta2, eps, wd, True, False) + p, m, v = params.clone(), torch.zeros(n), torch.zeros(n) + for step, grad in enumerate(grads, 1): + if in_kernel: + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, scale) + else: + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad.float() * scale, m, v) + results.append((p, m, v)) + op.destroy_adam(opt_id) + for name, a, b in zip(("params", "exp_avg", "exp_avg_sq"), *results): + assert torch.equal(a, b), f"{grad_dtype} {name}" + + +@pytest.mark.parametrize("earlier_steps", [0, 3], ids=["first_step", "later_step"]) +@pytest.mark.parametrize("grad_dtype", [torch.bfloat16, torch.float], ids=["bf16", "fp32"]) +def test_cpu_adam_rescale_step(earlier_steps, grad_dtype): + """adam_rescale_step turns a step on grad * old_scale into the step on grad * new_scale from the same state, up to + rounding: 7e-6 of the update for the params and 1e-6 of the largest moment (2.5e-7 and 1e-6 on the first step, + where the moments of before were 0). Rescaling to the same scale leaves the moments as they are, and takes the + update out of the params and adds it back.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + op = CPUAdamBuilder().load() + opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + lr, beta1, beta2, eps, wd, old_scale, new_scale = 1e-3, 0.9, 0.999, 1e-8, 0.01, 1.0, 0.2 + n = 100003 # not a multiple of the SIMD block, so the scalar tail runs too + torch.manual_seed(0) + params = torch.randn(n) * 0.02 + earlier = [torch.randn(n) * 1e-3 for _ in range(earlier_steps)] + grad = (torch.randn(n) * 1e-3).to(grad_dtype) + step = earlier_steps + 1 + results = [] + for rescale in (True, False): + op.create_adam(opt_id, lr, beta1, beta2, eps, wd, True, False) + p, m, v = params.clone(), torch.zeros(n), torch.zeros(n) + for t, g in enumerate(earlier, 1): + op.adam_update(opt_id, t, lr, beta1, beta2, eps, wd, True, p, g, m, v) + before = p.clone() + if rescale: + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, old_scale) + stepped = (p.clone(), m.clone(), v.clone()) + op.adam_rescale_step(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, old_scale, old_scale) + assert torch.equal(m, stepped[1]) and torch.equal(v, stepped[2]) + torch.testing.assert_close(p, stepped[0], rtol=0, atol=1e-5 * (stepped[0] - before).abs().max().item()) + op.adam_rescale_step(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, old_scale, new_scale) + else: + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, new_scale) + results.append((p, m, v)) + op.destroy_adam(opt_id) + (p, m, v), (p_ref, m_ref, v_ref) = results + torch.testing.assert_close(p, p_ref, rtol=0, atol=1e-4 * (p_ref - before).abs().max().item()) + torch.testing.assert_close(m, m_ref, rtol=0, atol=2e-6 * m_ref.abs().max().item()) + torch.testing.assert_close(v, v_ref, rtol=0, atol=1e-5 * v_ref.abs().max().item()) + + +def test_cpu_adam_rescale_step_rejects_l2_weight_decay(): + """Adam's L2 decay adds the params of before the step to the gradient, which a rescale cannot recover.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + op = CPUAdamBuilder().load() + opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + op.create_adam(opt_id, 1e-3, 0.9, 0.999, 1e-8, 0.01, False, False) + try: + p, g, m, v = torch.randn(64), torch.randn(64), torch.zeros(64), torch.zeros(64) + op.adam_update(opt_id, 1, 1e-3, 0.9, 0.999, 1e-8, 0.01, True, p, g, m, v) + with pytest.raises(RuntimeError, match="AdamW only"): + op.adam_rescale_step(opt_id, 1, 1e-3, 0.9, 0.999, 1e-8, 0.01, True, p, g, m, v, 1.0, 0.5) + finally: + op.destroy_adam(opt_id) + + @pytest.mark.parametrize('dtype', [torch.half, torch.bfloat16, torch.float], ids=["fp16", "bf16", "fp32"]) @pytest.mark.parametrize('model_size', [ From b2b46ae935c4a4ccdae5f1a6ac2c0321dfd22724 Mon Sep 17 00:00:00 2001 From: Xinyu Lian Date: Sun, 8 Nov 2026 02:24:29 +0000 Subject: [PATCH 10/10] SuperOffload: change the CPU steps of a clipped step in one pass instead 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 --- .../superoffload/superoffload_stage3.py | 75 ++++++++++--------- .../superoffload/superoffload_utils.py | 18 +++-- deepspeed/runtime/zero/stage3.py | 7 +- tests/unit/v1/zero/test_superoffload.py | 43 ++++++++++- 4 files changed, 95 insertions(+), 48 deletions(-) diff --git a/deepspeed/runtime/superoffload/superoffload_stage3.py b/deepspeed/runtime/superoffload/superoffload_stage3.py index 8b2bc1c3d084..adce271f234c 100644 --- a/deepspeed/runtime/superoffload/superoffload_stage3.py +++ b/deepspeed/runtime/superoffload/superoffload_stage3.py @@ -47,7 +47,7 @@ def _cpu_adam_settings(optimizer): class _CPUStep(NamedTuple): - """A CPU step submitted during backward, and what a rollback of it needs.""" + """A CPU step submitted during backward, and what changing or undoing it needs.""" # Its result tells whether the sub-group stepped: one whose gradients were not finite is left to step(). stepped: Future hyperparams: dict @@ -89,8 +89,9 @@ def __init__( super().__init__(module, init_optimizer, param_names, timers, ds_config, **kwargs) self.sub_group_to_param_num = [len(group) for group in self.fp16_groups] - # A step taken during backward is undone with adam_rollback, which supports only fp32 AdamW with bias - # correction. With other settings, step() runs the CPU sub-groups after the overflow and clipping checks. + # step() changes a step taken during backward with adam_rescale_step when the gradients clip or the loss scale + # changed, and undoes it with adam_rollback when it skips the step. Both need fp32 AdamW, and adam_rollback + # bias correction. With other settings, step() runs the CPU sub-groups after the overflow and clipping checks. bias_correction = all(group.get("bias_correction", True) for group in self.optimizer.param_groups) self.speculative_cpu_step = (adamw_mode and fp32_optimizer_states and self.master_weights_and_grads_dtype == torch.float32 and bias_correction) @@ -160,17 +161,6 @@ def _worker_grad(self, sub_group_id): half_grad = self._half_cpu_grads.get(sub_group_id) return half_grad if half_grad is not None else self.fp32_partitioned_groups_flat[sub_group_id].grad.data - def _cpu_step_grad(self, sub_group_id, total_norm): - """The gradient for step() to step a CPU sub-group with that did not step during backward, unscaled and - clipped. A bf16 gradient that needs neither goes as it is; otherwise the fp32 gradient takes the scaling.""" - half_grad = self._half_cpu_grads.get(sub_group_id) - if half_grad is not None: - if self.loss_scale == 1 and not self.check_clip_grads(total_norm): - return half_grad - self.fp32_partitioned_groups_flat[sub_group_id].grad.copy_(half_grad) - self.unscale_and_clip_grads(sub_group_id, total_norm) - return self.fp32_partitioned_groups_flat[sub_group_id].grad.data - def destroy(self): self.cpu_adam.close() super().destroy() @@ -463,22 +453,20 @@ def step(self, closure=None): timer_names.add(OPTIMIZER_STEP_TIMER) self.timers(OPTIMIZER_STEP_TIMER).start() - # The steps taken during backward used the unclipped gradients, and the loss scale of the backward, which - # _loss_scale_update_and_overflow_cleanup may have raised before unscale_and_clip_grads removes it. - if self.check_clip_grads(scaled_global_grad_norm) or self._loss_scale_changed(): - self._rollback_submitted_cpu_sub_groups() - + # CPU Adam removes the loss scale and clips as it reads the gradients, by the factor unscale_and_clip_grads + # multiplies them by. + grad_scale = float(self._unscale_and_clip_factor(scaled_global_grad_norm)) for sub_group_id, group in enumerate(self.fp16_groups): self._prepare_sub_group(sub_group_id, timer_names) - # A sub-group stepped during backward had the loss scale removed in the worker, and one that needed - # clipping was rolled back above, so its gradients are not read again. if self.subgroup_to_device[sub_group_id] != 'cpu': self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) self._optimizer_step(sub_group_id) - elif sub_group_id not in self._submitted_cpu_sub_groups: + elif sub_group_id in self._submitted_cpu_sub_groups: + self._rescale_cpu_step(sub_group_id, grad_scale) + else: param_group_id = self.sub_group_to_group_id[sub_group_id] - self._cpu_step_and_wait(sub_group_id, self._cpu_step_grad(sub_group_id, scaled_global_grad_norm), - self._hyperparams(param_group_id)) + self._cpu_step_and_wait(sub_group_id, self._worker_grad(sub_group_id), + self._hyperparams(param_group_id), grad_scale) if sub_group_id in self._staged_lp_params: self._apply_staged_lp_params(sub_group_id) @@ -544,10 +532,6 @@ def _hyperparams(self, param_group_id): hyperparams["bias_correction"] = group.get("bias_correction", True) return hyperparams - def _loss_scale_changed(self): - grad_scale = 1.0 / self.loss_scale - return any(step.grad_scale != grad_scale for step in self._submitted_cpu_sub_groups.values()) - def _step_when_copied(self, sub_group_id, copies_done, grads_finite, hyperparams, grad_scale): """In the CPU Adam thread: step a sub-group once its gradients are on the CPU, and tell whether it did. @@ -570,8 +554,33 @@ def _cpu_step_and_wait(self, sub_group_id, grad, hyperparams, grad_scale=1.0, ro self._drop_staged_lp_params(sub_group_id) self.cpu_adam.submit(self._cpu_adam_step, sub_group_id, grad, hyperparams, grad_scale, rollback).result() + def _rescale_cpu_step(self, sub_group_id, grad_scale): + """Turn the CPU step a sub-group took during backward into the step on its gradients at grad_scale. + + That step read the gradients unclipped and with the loss scale of the backward, which + _loss_scale_update_and_overflow_cleanup may have raised since. One pass over the weights and moments changes + the step, where undoing it and stepping again would take two, and the undo is slower than a step. + """ + step = self._submitted_cpu_sub_groups[sub_group_id] + if step.grad_scale == grad_scale: + return + hyperparams = self._hyperparams(self.sub_group_to_group_id[sub_group_id]) + if hyperparams != step.hyperparams: + # Changed since backward, so the step taken with the old ones is undone, and taken again. + self._rollback_cpu_sub_group(sub_group_id) + self._cpu_step_and_wait(sub_group_id, self._worker_grad(sub_group_id), hyperparams, grad_scale) + return + self._drop_staged_lp_params(sub_group_id) + self.cpu_adam.submit(self._cpu_adam_rescale, sub_group_id, hyperparams, step.grad_scale, grad_scale).result() + + def _cpu_adam_rescale(self, sub_group_id, hyperparams, old_scale, new_scale): + """In the CPU Adam thread: turn the last step of a CPU sub-group into the step on its gradient at new_scale.""" + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + self.cpu_adam.rescale_step(self.optimizer.state[fp32_param], fp32_param.data, self._worker_grad(sub_group_id), + hyperparams, old_scale, new_scale) + def _rollback_submitted_cpu_sub_groups(self): - """Undo the CPU steps taken during backward, for a skipped step or before redoing them with clipping.""" + """Undo the CPU steps taken during backward, for a skipped step.""" for sub_group_id in list(self._submitted_cpu_sub_groups): self._rollback_cpu_sub_group(sub_group_id) @@ -584,14 +593,6 @@ def _rollback_cpu_sub_group(self, sub_group_id): step.grad_scale, rollback=True) - @instrument_w_nvtx - def check_clip_grads(self, total_norm): - """Check if unscale_and_clip_grads will scale the gradients down, with the same test.""" - if not self.clip_grad: - return False - clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad - return bool(clip > 1.0) - def load_state_dict(self, state_dict_list, load_optimizer_states=True, diff --git a/deepspeed/runtime/superoffload/superoffload_utils.py b/deepspeed/runtime/superoffload/superoffload_utils.py index d036ab7dd495..143363654860 100644 --- a/deepspeed/runtime/superoffload/superoffload_utils.py +++ b/deepspeed/runtime/superoffload/superoffload_utils.py @@ -66,7 +66,7 @@ def __init__(self, defaults: Dict[str, Any], adamw_mode: bool, fp32_optimizer_st DeepSpeedCPUAdam.optimizer_id += 1 self._ops.create_adam(self._opt_id, defaults["lr"], defaults["betas"][0], defaults["betas"][1], defaults["eps"], defaults["weight_decay"], adamw_mode, should_log_le("info")) - # Holds a gradient with the loss scale removed, or one in fp32 for adam_rollback. + # Holds the gradient adam_rollback reads: in fp32, with the loss scale removed. self._grad_buffer = None def submit(self, fn: Callable, *args) -> Future: @@ -86,11 +86,17 @@ def step(self, state["exp_avg"] = torch.zeros_like(weights, dtype=state_dtype) state["exp_avg_sq"] = torch.zeros_like(weights, dtype=state_dtype) state["step"] += 1 - if grad_scale != 1.0: - grad = self._unscaled(grad, weights.dtype, grad_scale) - # adam_update reads bf16 gradients of fp32 weights as they are, so they go to the CPU at half the bytes. + # adam_update reads bf16 gradients of fp32 weights as they are, so they go to the CPU at half the bytes, and + # multiplies them by grad_scale as it reads them. self._ops.adam_update(self._opt_id, state["step"], *_kernel_args(hyperparams), weights, grad, state["exp_avg"], - state["exp_avg_sq"]) + state["exp_avg_sq"], grad_scale) + + def rescale_step(self, state: Dict[str, Any], weights: torch.Tensor, grad: torch.Tensor, + hyperparams: Dict[str, Any], old_scale: float, new_scale: float) -> None: + """Turn the last step of the weights, taken on grad * old_scale with these hyperparameters, into the step on + grad * new_scale, in one pass over the weights and moments.""" + self._ops.adam_rescale_step(self._opt_id, state["step"], *_kernel_args(hyperparams), weights, grad, + state["exp_avg"], state["exp_avg_sq"], old_scale, new_scale) def rollback(self, state: Dict[str, Any], @@ -111,7 +117,7 @@ def _unscaled(self, grad: torch.Tensor, dtype: torch.dtype, grad_scale: float) - if buffer is None or buffer.dtype != dtype or buffer.numel() < grad.numel(): buffer = self._grad_buffer = torch.empty(grad.numel(), dtype=dtype) unscaled = buffer[:grad.numel()].view_as(grad) - # A step submitted during backward gets the gradient before the loss scale is removed. + # The step of a sub-group submitted during backward read the gradient before the loss scale is removed. torch.mul(grad, grad_scale, out=unscaled) return unscaled diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 563f5c4c0050..e3246d4fc618 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -3036,6 +3036,10 @@ def dump_post_step_gradients(self): @instrument_w_nvtx def unscale_and_clip_grads(self, sub_group_id, total_norm): + self.fp32_partitioned_groups_flat[sub_group_id].grad.mul_(self._unscale_and_clip_factor(total_norm)) + + def _unscale_and_clip_factor(self, total_norm): + """The factor unscale_and_clip_grads multiplies the gradients by: it removes the loss scale and clips.""" # compute combined scale factor for this group combined_scale = self.loss_scale if self.clip_grad > 0.: @@ -3043,8 +3047,7 @@ def unscale_and_clip_grads(self, sub_group_id, total_norm): clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad clip = torch.clamp(clip, min=1.0) combined_scale = clip * self.loss_scale - - self.fp32_partitioned_groups_flat[sub_group_id].grad.mul_(1. / combined_scale) + return 1. / combined_scale def _check_overflow(self, partition_gradients=True): self.overflow = self.has_overflow(partition_gradients) diff --git a/tests/unit/v1/zero/test_superoffload.py b/tests/unit/v1/zero/test_superoffload.py index 0a7f0cd67bd1..c88c3f50412e 100644 --- a/tests/unit/v1/zero/test_superoffload.py +++ b/tests/unit/v1/zero/test_superoffload.py @@ -179,8 +179,8 @@ def _b_and_c_in_one_micro_step(unused): "bf16_master_weights": (dict(precision="bf16_master"), lambda unused: _all_used, True), "bf16_optimizer_states": (dict(precision="bf16_states"), lambda unused: _all_used, True), "two_param_groups": (dict(two_groups=True, sub_group_size=4500), _b_unused_on_even_steps, True), - # Clipping undoes the steps taken during backward and redoes them, which is exact up to rounding. In fp32 the - # rounding stays small; bf16 weights would amplify it. + # Clipping turns the steps taken during backward into steps on the clipped gradients, exact up to rounding. In + # fp32 the rounding stays small; bf16 weights would amplify it. "clipping": (dict(clip=0.05, precision="fp32"), lambda unused: _all_used, False), # A step cannot be undone with a beta of 0 (the rollback divides by the betas), so those CPU steps wait for # step() and stay exact. @@ -253,7 +253,7 @@ def test_large_grads(self): def test_fp16_loss_scale_increase(self): """Dynamic loss scaling raises the scale in step(), after the CPU steps taken during backward used the old - one; those steps are undone and redone, exact up to rounding, which fp16 weights amplify (1.4e-5 after 6 + one; those steps are changed to the new one, exact up to rounding, which fp16 weights amplify (1.4e-5 after 6 steps; 3.9e-3 when the steps are kept).""" results = [] for super_offload in (False, True): @@ -295,6 +295,43 @@ def record(state, weights, grad, *args): assert not any(t.is_shared() for t in engine.optimizer.fp32_partitioned_groups_flat) _close(engine) + def test_clipping_bf16(self): + """Clipping with bf16 gradients, which the worker scales as it reads them. Where a master differs by an ulp, + its bf16 weight can round the other way, which the next steps amplify: 3.5e-7 after 2 steps, 6e-5 after 8 + (3e-5 when the steps of backward are undone and taken again), so only 2 steps are compared.""" + results = [] + for super_offload in (False, True): + engine = _engine(_config(super_offload=super_offload, clip=0.05)) + results.append(_train(engine, 2, _all_used)) + _close(engine) + _assert_same(results[1], results[0], exact=False) + + def test_clipped_step_rescales_cpu_steps(self): + """A step that clips changes each CPU step taken during backward to the clipped gradients in one pass, + instead of undoing it and stepping again.""" + engine = _engine(_config(super_offload=True, clip=0.05)) + cpu_adam = engine.optimizer.cpu_adam + calls = [] + + def recording(name): + update = getattr(cpu_adam, name) + + def record(*args): + calls.append(name) + update(*args) + + return record + + for name in ("step", "rescale_step", "rollback"): + setattr(cpu_adam, name, recording(name)) + x = torch.randn(4, HIDDEN).to(engine.device, torch.bfloat16) + engine.backward(engine(x, {})) + engine.step() + assert engine.get_global_grad_norm() > 0.05 + cpu_sub_groups = list(engine.optimizer.subgroup_to_device.values()).count("cpu") + assert calls == ["step"] * cpu_sub_groups + ["rescale_step"] * cpu_sub_groups + _close(engine) + def test_cpu_buffers_pinned(self): config = _config(super_offload=True) config["zero_optimization"]["offload_optimizer"]["pin_memory"] = True