# Copyright (c) DeepSpeed Team. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team """ SuperOffload utilities for 1) running CPU optimizers in separate processes. """ from typing import Dict, Optional, Any import torch import torch.multiprocessing as mp import psutil from deepspeed.ops.adam import DeepSpeedCPUAdam from deepspeed.utils import logger class TaskKeys: PARAM_DATA = "param_data" PARAM_GRAD = "param_grad" PARAM_GROUP_ID = "param_group_id" SUB_GROUP_ID = "sub_group_id" ROLLBACK = "rollback" LR = "lr" class ResultKeys: UPDATED_PARAM = "updated_param" EVENT_TYPE = "event_type" class EventTypes: ADAM_STEP = "adam_step" ROLLBACK = "rollback" def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, result_queue: mp.SimpleQueue, optimizer_config: Dict[str, Any], max_grad_numel: int) -> 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, and amsgrad parameters max_grad_numel: Maximum number of elements expected in gradient tensors """ 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"], betas=first_cfg["betas"], eps=first_cfg["eps"], weight_decay=first_cfg["weight_decay"], amsgrad=first_cfg["amsgrad"]) 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"], "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}) return # Pre-allocate reusable pinned memory buffer for gradients pinned_grad_buffer = torch.empty(max_grad_numel, dtype=torch.float32, device='cpu', pin_memory=True) while True: try: task = param_queue.get() if task is None: logger.debug("Received termination signal, shutting down worker") break 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) 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 task_lr is not None: optimizer.param_groups[param_group_id]['lr'] = task_lr 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 = pinned_grad_buffer[:grad_numel].view_as(param_grad) param_grad_cpu.copy_(param_grad, non_blocking=False) 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 pinned memory buffer if 'pinned_grad_buffer' in locals(): del pinned_grad_buffer logger.debug("Cleaned up pinned memory buffer") logger.debug("Worker process terminated") class SuperOffloadCPUOptimizer: def __init__(self, optimizer_config: Dict[str, Any], cpuadam_cores_perc: float = 0.8, max_grad_numel: int = 1000000) -> None: 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), 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, lr: float = None) -> None: """ Queue parameter for optimization in the worker process. """ 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 lr is not None: task[TaskKeys.LR] = lr self.param_queue.put(task) 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 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")