[Feature] offload optimizer states to CPU (reduce NPU memory, trade-off performance)

Author: wentiange

xtuner/v1/config/optim.py

@@ -9,11 +9,10 @@
 
 from xtuner.v1.optim import Muon
 from xtuner.v1.utils import get_logger
-
+import types
 
 logger = get_logger()
 
-
 class OptimConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
     lr: Annotated[float, Parameter(help="Learning rate for optimization")] = 1e-5
@@ -32,6 +31,7 @@ class AdamWConfig(OptimConfig):
     betas: Annotated[Tuple[float, float], Parameter(help="Beta coefficients for Adam optimizer")] = (0.9, 0.95)
     eps: Annotated[float, Parameter(help="Epsilon value for numerical stability in Adam optimizer")] = 1e-8
     foreach: Annotated[Optional[bool], Parameter(help="Use foreach implementation for AdamW")] = None
+    swap_optimizer: Annotated[Optional[bool], Parameter(help="Swap optimizer states to host memory.")] = False
 
     def build(self, model):
         params = [p for p in model.parameters() if p.requires_grad]
@@ -52,10 +52,13 @@ def build(self, model):
                 f"Total trainable parameters: {num_total_requires_grad // 1e6}M, total parameters: {num_total // 1e6}M"
             )
             logger.info(f"Untrainable parameters names: {untrainable_names}")
-        return torch.optim.AdamW(
+        optimizer = torch.optim.AdamW(
             params, lr=self.lr, betas=self.betas, eps=self.eps, weight_decay=self.weight_decay, foreach=self.foreach
-        )
-
+            )
+        if self.swap_optimizer:
+            SwapOptimizerOperate(optimizer).opt_states_initialization()
+            optimizer.step = types.MethodType(swap_adamw_step, optimizer)
+        return optimizer
 
 class MuonConfig(OptimConfig):
     weight_decay: Annotated[float, Parameter(help="Weight decay coefficient for L2 regularization")] = 0.1
@@ -134,3 +137,156 @@ class LRConfig(BaseModel):
     )
     warmup_ratio: Annotated[float, Parameter(help="Ratio of warmup steps to total training steps")] = 0.03
     lr_min: Annotated[float, Parameter(help="Minimum learning rate for optimization")] = 1e-6
+
+class SwapOptimizerOperate():
+
+    swap_to_device_stream = None
+    swap_to_host_stream = None
+
+    swap_to_device_events_map = {}
+    swap_to_host_events_map = {}
+
+    param_to_cpu_states_map = {}
+    param_to_device_states_map = {}
+
+    state_keys = ['exp_avg', 'exp_avg_sq', 'max_exp_avg_sq']
+
+    def __init__(self, optimizer, swap_optimizer_times=16):
+        self.optimizer = optimizer
+        self.swap_optimizer_times = swap_optimizer_times
+        if SwapOptimizerOperate.swap_to_device_stream is None:
+            SwapOptimizerOperate.swap_to_device_stream = torch.npu.Stream()
+            SwapOptimizerOperate.swap_to_host_stream = torch.npu.Stream()
+
+        # create all parameters list for step
+        self.optimizer.param_to_group_map = {}
+
+        for group in self.optimizer.param_groups:
+            for p in group['params']:
+                self.optimizer.param_to_group_map[p] = group
+
+        # print swap param num and size
+        swap_num = sum([main_param.to_local().numel() for main_param in self.optimizer.param_to_group_map])
+        swap_numel = swap_num // self.swap_optimizer_times
+        self.optimizer.swap_numel = swap_numel
+    
+        swap_memory = swap_num * 8 / 1024 / 1024
+        print('[Rank {}] swap optimizer param num: {},  param size: {}MB\n'.format(torch.npu.current_device(), swap_num, swap_memory), end='')
+
+    def opt_states_initialization(self):
+        for group in self.optimizer.param_groups:
+            for param in group["params"]:
+                device_state_dtensor = self.optimizer.state[param]
+                device_state_tensor = {}
+                cpu_state = {}
+                
+                amsgrad = self.optimizer.param_to_group_map[param]['amsgrad']
+
+                for key in self.state_keys:
+                    if key == 'max_exp_avg_sq' and not amsgrad:
+                        device_state_dtensor[key] = None
+                        device_state_tensor[key] = None
+                        cpu_state[key] = None
+                    else:
+                        device_state_dtensor[key] = torch.zeros_like(param, memory_format=torch.preserve_format)
+                        # convert dtensor to tensor
+                        device_state_tensor[key] = device_state_dtensor[key].to_local()
+                        
+                        cpu_state[key] = torch.empty_like(device_state_tensor[key], pin_memory=True, device='cpu')
+                        cpu_state[key].copy_(device_state_tensor[key], non_blocking=True)
+
+                        device_state_tensor[key].storage().resize_(0)
+
+                self.param_to_device_states_map[param] = device_state_tensor
+                self.param_to_cpu_states_map[param] = cpu_state
+        torch.npu.synchronize()
+
+    @classmethod
+    def swap_all_to_host(cls):
+        for param in cls.param_to_cpu_states_map.keys():
+            cls.swap_tensors_to_host(param)
+        for param in cls.param_to_cpu_states_map.keys():
+            event = cls.swap_to_host_events_map.get(param, None)
+            if event is not None:
+                torch.npu.current_stream().wait_event(event)
+                cls.swap_to_host_events_map[param] = None     
+
+    @classmethod
+    def swap_all_to_device(cls):
+        for param in cls.param_to_cpu_states_map.keys():
+            cls.swap_tensors_to_device(param)
+        for param in cls.param_to_cpu_states_map.keys():
+            event = cls.swap_to_device_events_map.get(param, None)
+            if event is not None:
+                torch.npu.current_stream().wait_event(event)
+                cls.swap_to_device_events_map[param] = None     
+
+    @classmethod
+    def swap_tensors_to_device(cls, param):
+        cpu_state = cls.param_to_cpu_states_map[param]
+
+        if param in cls.param_to_device_states_map:
+            device_state = cls.param_to_device_states_map[param]
+            for key in cls.state_keys:
+                if device_state[key] is not None and device_state[key].storage().size() == 0:
+                    device_state[key].storage().resize_(cpu_state[key].storage().size())
+                    device_state[key].copy_(cpu_state[key], non_blocking=True)
+
+        cls.swap_to_device_events_map[param] =  torch.npu.current_stream().record_event()
+
+    @classmethod
+    def wait_swap_to_device_event(cls, param):
+        event = cls.swap_to_device_events_map.get(param, None)
+        if event is not None:
+            torch.npu.current_stream().wait_event(event)
+            cls.swap_to_device_events_map[param] = None
+
+    @classmethod
+    def swap_tensors_to_host(cls, param):
+        cpu_state = cls.param_to_cpu_states_map[param]
+
+        if param in cls.param_to_device_states_map:
+            device_state = cls.param_to_device_states_map[param]
+            for key in cls.state_keys:
+                if key in device_state and device_state[key] is not None and device_state[key].storage().size() != 0:
+                    cpu_state[key].copy_(device_state[key], non_blocking=True)
+                    device_state[key].storage().resize_(0)
+
+        cls.swap_to_host_events_map[param] = torch.npu.current_stream().record_event()
+
+def swap_adamw_step(self, closure=None):
+    loss = None
+    if closure is not None:
+        with torch.enable_grad():
+            loss = closure()
+
+    for group in self.param_groups:
+        if 'step' in group:
+            group['step'] += 1
+            if group['step'].is_cpu:
+                group['step'] = group['step'].npu()
+        else:
+            group['step'] = torch.tensor(1, dtype=torch.int64, device=torch.npu.current_device())
+    
+    params_list = list(self.param_to_group_map.keys())
+
+    SwapOptimizerOperate.swap_all_to_device()
+
+    for i, param in enumerate(params_list):
+        if param.grad is None:
+            continue
+        if param.grad.is_sparse:
+            raise RuntimeError('AdamW does not support sparse gradients')
+
+        group = self.param_to_group_map[param]
+        amsgrad = group['amsgrad']
+        beta1, beta2 = group['betas']
+        state = self.state[param]
+
+        torch._fused_adamw_([param.to_local()], [param.grad.to_local()], [state['exp_avg'].to_local()], [state['exp_avg_sq'].to_local()], [state['max_exp_avg_sq']] if amsgrad else [],
+                                 [group['step']], amsgrad=amsgrad, lr=group['lr'], beta1=beta1, beta2=beta2, weight_decay=group['weight_decay'],
+                                 eps=group['eps'], maximize=group['maximize'])
+
+    # it maybe removed 
+    torch.npu.synchronize()
+    return loss

xtuner/v1/engine/train_engine.py

@@ -511,7 +511,7 @@ def put_model_to_device(self, device: torch.device | str):
 
     def put_optimizer_to_device(self, device: torch.device | str):
         """Put the optimizer to the given device."""
-        if self.fsdp_cfg.cpu_offload:
+        if self.fsdp_cfg.cpu_offload or self.optim_cfg.swap_optimizer:
             return
         if not self.optimizer.state:
             return

xtuner/v1/train/trainer.py

@@ -737,6 +737,13 @@ def fit(self):
             grad_norm = self._engine.clip_grad_norm(do_clip=self._do_clip, dtype=self._grad_norm_dtype)
             self._engine.step_optimizer(grad_norm)
 
+            if self._optim_config.swap_optimizer:
+                is_save_dcp = (self._is_save_dcp(is_snapshot=False) or self._is_save_dcp(is_snapshot=True))
+                if not is_save_dcp:
+                    from xtuner.v1.config.optim import SwapOptimizerOperate
+                    SwapOptimizerOperate.swap_all_to_host()
+                torch.npu.synchronize()
+
             time_after_train_step = time.time()
             ProberList.after_step()
 
@@ -1092,6 +1099,22 @@ def _maybe_check_health(self):
                 raise RuntimeError("Health check failed, exit training")
             logger.info(f"Health check passed at step {self.cur_step}")
 
+    
+    def _is_save_dcp(self, is_snapshot: bool = False) -> bool:
+        ckp_interval = self._checkpoint_interval if not is_snapshot else self._snapshot_interval
+        cur_step = self._cur_step + 1
+        if ckp_interval is None:
+            return False
+
+        if ckp_interval == -1:  # only save at the end of training
+            if cur_step != self.total_step:
+                return False
+        else:
+            if cur_step % ckp_interval != 0 and (is_snapshot or cur_step != self.total_step):
+                # if is_snapshot, only save at interval
+                # else save at interval or at the end of training
+                return False
+        return True
     def _maybe_save(self, is_snapshot: bool = False) -> bool:
         ckp_interval = self._checkpoint_interval if not is_snapshot else self._snapshot_interval
         if ckp_interval is None:
@@ -1123,6 +1146,12 @@ def _maybe_save(self, is_snapshot: bool = False) -> bool:
             optimizer_dir=optimizer_path,
         )
 
+        if self._optim_config.swap_optimizer:
+            torch.npu.synchronize()
+            from xtuner.v1.config.optim import SwapOptimizerOperate
+            SwapOptimizerOperate.swap_all_to_host()
+            torch.npu.synchronize()
+
         # Save dataloader
         self._save_dataloader(dataloader_path)
 
@@ -1759,6 +1788,11 @@ def _load_checkpoint(self):
             else None
         )
 
+        if self._optim_config.swap_optimizer:
+            from xtuner.v1.config.optim import SwapOptimizerOperate
+            SwapOptimizerOperate.swap_all_to_device()
+
+
         self._engine.load_dcp(
             model_dir=model_path,
             optimizer_dir=optimizer_path,