perf(fsdp): accelerate pinned sleep and wake

miles/backends/fsdp_utils/actor.py

@@ -263,8 +263,11 @@ def sleep(self) -> None:
 
         print_memory("before offload DiT")
 
-        self.model.cpu()
-        move_torch_optimizer(self.optimizer, "cpu")
+        move_torch_model(self.model, "cpu", pin_memory=True)
+        move_torch_optimizer(self.optimizer, "cpu", pin_memory=True)
+        if self.ema_shadow is not None:
+            self.ema_shadow.to("cpu", non_blocking=True, pin_memory=True)
+        torch.cuda.synchronize()
         clear_memory()
         dist.barrier(group=get_gloo_group())
         print_memory("after sleep DiT")
@@ -274,8 +277,12 @@ def wake_up(self) -> None:
         if not self.args.offload_train:
             return
 
-        self.model.cuda()
-        move_torch_optimizer(self.optimizer, "cuda")
+        device = torch.device("cuda", torch.cuda.current_device())
+        move_torch_model(self.model, device)
+        move_torch_optimizer(self.optimizer, device)
+        if self.ema_shadow is not None:
+            self.ema_shadow.to(device, non_blocking=True)
+        torch.cuda.synchronize()
         dist.barrier(group=get_gloo_group())
         print_memory("after wake_up DiT")
 
@@ -573,20 +580,123 @@ def _compute_noise_pred() -> torch.Tensor:
         )
 
 
+def _to_pinned_cpu(tensor: torch.Tensor) -> torch.Tensor:
+    if isinstance(tensor, DTensor):
+        local_tensor = tensor.to_local()
+        if local_tensor.device.type == "cpu" and local_tensor.is_pinned():
+            return tensor
+        pinned = torch.empty_like(local_tensor, device="cpu", pin_memory=True)
+        pinned.copy_(local_tensor, non_blocking=local_tensor.device.type == "cuda")
+        return DTensor(pinned, tensor._spec, requires_grad=tensor.requires_grad)
+    if tensor.device.type == "cpu" and tensor.is_pinned():
+        return tensor
+    pinned = torch.empty_like(tensor, device="cpu", pin_memory=True)
+    pinned.copy_(tensor, non_blocking=tensor.device.type == "cuda")
+    return pinned
+
+
+def _to_device(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
+    if isinstance(tensor, DTensor):
+        local_tensor = tensor.to_local()
+        if local_tensor.device == device:
+            return tensor
+        moved = local_tensor.to(device, non_blocking=True)
+        return DTensor(moved, tensor._spec, requires_grad=tensor.requires_grad)
+    return tensor.to(device, non_blocking=True)
+
+
+def _fsdp_modules(model: torch.nn.Module):
+    from torch.distributed.fsdp import FSDPModule
+
+    fsdp_modules = [module for module in model.modules() if isinstance(module, FSDPModule)]
+    for module in fsdp_modules:
+        module.reshard()
+    return fsdp_modules
+
+
+def _fsdp_params(fsdp_modules):
+    fsdp_params = []
+    for module in fsdp_modules:
+        state = module._get_fsdp_state()
+        if state._fsdp_param_group is not None:
+            fsdp_params.extend(state._fsdp_param_group.fsdp_params)
+    return fsdp_params
+
+
+def _set_fsdp_param_storage(fsdp_param, storage: torch.Tensor) -> None:
+    sharded_param = fsdp_param.sharded_param
+    shard_dim = fsdp_param.fsdp_placement.dim
+    padded = storage.view(fsdp_param.padded_sharded_param_size)
+    local = padded.narrow(shard_dim, 0, fsdp_param.sharded_size[shard_dim])
+    replacement = torch.nn.Parameter(
+        fsdp_param.to_sharded_dtensor(local),
+        requires_grad=sharded_param.requires_grad,
+    )
+    torch.utils.swap_tensors(sharded_param, replacement)
+    fsdp_param._sharded_param_data = storage
+
+
+def _move_model_to_pinned_cpu(model: torch.nn.Module) -> None:
+    fsdp_modules = _fsdp_modules(model)
+    fsdp_params = _fsdp_params(fsdp_modules)
+    managed_param_ids = {id(fsdp_param.sharded_param) for fsdp_param in fsdp_params}
+    for parameter in model.parameters():
+        parameter.grad = None
+        if id(parameter) not in managed_param_ids:
+            parameter.data = _to_pinned_cpu(parameter.data)
+
+    for fsdp_param in fsdp_params:
+        _set_fsdp_param_storage(fsdp_param, _to_pinned_cpu(fsdp_param._sharded_param_data))
+
+    for module in model.modules():
+        for name, buffer in module._buffers.items():
+            if buffer is not None:
+                module._buffers[name] = _to_pinned_cpu(buffer)
+
+
+def _move_model_to_device(model: torch.nn.Module, device: torch.device) -> None:
+    fsdp_modules = _fsdp_modules(model)
+    fsdp_params = _fsdp_params(fsdp_modules)
+    managed_param_ids = {id(fsdp_param.sharded_param) for fsdp_param in fsdp_params}
+    for parameter in model.parameters():
+        if id(parameter) not in managed_param_ids:
+            parameter.data = _to_device(parameter.data, device)
+
+    for fsdp_param in fsdp_params:
+        _set_fsdp_param_storage(fsdp_param, _to_device(fsdp_param._sharded_param_data, device))
+
+    for module in model.modules():
+        for name, buffer in module._buffers.items():
+            if buffer is not None:
+                module._buffers[name] = _to_device(buffer, device)
+
+
 @torch.no_grad()
-def move_torch_optimizer(optimizer, device):
+def move_torch_model(model: torch.nn.Module, device, *, pin_memory: bool = False) -> None:
+    device = torch.device(device)
+    if pin_memory:
+        if device.type != "cpu":
+            raise ValueError("pin_memory requires a CPU destination")
+        _move_model_to_pinned_cpu(model)
+    else:
+        _move_model_to_device(model, device)
+
+
+@torch.no_grad()
+def move_torch_optimizer(optimizer, device, *, pin_memory: bool = False):
     """ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py"""
     if not optimizer.state:
         return
 
+    device = torch.device(device)
+    if pin_memory and device.type != "cpu":
+        raise ValueError("pin_memory requires a CPU destination")
     for param_group in optimizer.param_groups:
         for param in param_group["params"]:
             state = optimizer.state[param]
             for key, value in state.items():
                 if isinstance(value, torch.Tensor):
-                    state[key] = value.to(device, non_blocking=True)
-
-    torch.cuda.synchronize()
+                    state[key] = _to_pinned_cpu(value) if pin_memory else _to_device(value, device)
 
 
 def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) -> torch.nn.Module:

miles/backends/fsdp_utils/ema.py

@@ -54,6 +54,21 @@ def update(self) -> float:
             sh.mul_(delta).add_(_local(live.detach()).to(sh.device), alpha=1.0 - delta)
         return delta
 
+    @torch.no_grad()
+    def to(self, device, *, non_blocking: bool = False, pin_memory: bool = False) -> None:
+        device = torch.device(device)
+        if pin_memory and device.type != "cpu":
+            raise ValueError("pin_memory requires a CPU destination")
+        for index, shadow in enumerate(self.shadow):
+            if pin_memory:
+                if shadow.device.type == "cpu" and shadow.is_pinned():
+                    continue
+                moved = torch.empty_like(shadow, device="cpu", pin_memory=True)
+                moved.copy_(shadow, non_blocking=non_blocking and shadow.device.type == "cuda")
+                self.shadow[index] = moved
+            else:
+                self.shadow[index] = shadow.to(device, non_blocking=non_blocking)
+
     @contextmanager
     def swap_in(self):
         """Temporarily expose EMA weights as the live parameters."""

tests/fast-gpu/backends/fsdp_utils/_sleep_wakeup_worker.py

@@ -0,0 +1,192 @@
+import os
+import time
+
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import FSDPModule, fully_shard
+from torch.distributed.tensor import DTensor
+
+from miles.backends.fsdp_utils.actor import move_torch_model, move_torch_optimizer
+from miles.backends.fsdp_utils.ema import EmaShadow
+
+
+class _Model(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = torch.nn.ModuleList([torch.nn.Linear(8, 7), torch.nn.Linear(7, 8)])
+        self.register_buffer("scale", torch.arange(8, dtype=torch.float32))
+        self.register_buffer("transposed", torch.arange(16, dtype=torch.float32).reshape(4, 4).T)
+
+    def forward(self, inputs):
+        for layer in self.layers:
+            inputs = torch.nn.functional.relu(layer(inputs))
+        return inputs * self.scale
+
+
+def _local(tensor):
+    return tensor.to_local() if isinstance(tensor, DTensor) else tensor
+
+
+def _assert_device(tensors, device_type, *, pinned=False):
+    for tensor in tensors:
+        local = _local(tensor)
+        assert local.device.type == device_type, (
+            f"expected {device_type}, got tensor={type(tensor).__name__} "
+            f"device={tensor.device} local_device={local.device}"
+        )
+        if device_type == "cpu":
+            assert local.is_pinned() is pinned
+
+
+def _optimizer_tensors(optimizer):
+    return [value for state in optimizer.state.values() for value in state.values() if isinstance(value, torch.Tensor)]
+
+
+def _fsdp_params(model):
+    params = []
+    for module in model.modules():
+        if not isinstance(module, FSDPModule):
+            continue
+        state = module._get_fsdp_state()
+        if state._fsdp_param_group is not None:
+            params.extend(state._fsdp_param_group.fsdp_params)
+    return params
+
+
+def _layout_signature(tensor):
+    local = _local(tensor)
+    dtensor_layout = (
+        (
+            tuple(tensor.placements),
+            tensor.device_mesh.device_type,
+            tuple(tensor.device_mesh.shape),
+            tensor.device_mesh.mesh_dim_names,
+            tuple(tensor.device_mesh.mesh.flatten().cpu().tolist()),
+        )
+        if isinstance(tensor, DTensor)
+        else None
+    )
+    return (
+        type(tensor),
+        tensor.layout,
+        tuple(tensor.shape),
+        tensor.stride(),
+        local.layout,
+        tuple(local.shape),
+        local.stride(),
+        local.storage_offset(),
+        dtensor_layout,
+    )
+
+
+def main():
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group("nccl")
+
+    device = torch.device("cuda", local_rank)
+    mesh = init_device_mesh("cuda", (dist.get_world_size(),))
+    model = _Model().to(device)
+    for layer in model.layers:
+        fully_shard(layer, mesh=mesh)
+    fully_shard(model, mesh=mesh)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
+    ema = EmaShadow(model.parameters(), decay=1.0, flat_steps=10)
+
+    inputs = torch.randn(4, 8, device=device)
+    model(inputs).sum().backward()
+    optimizer.step()
+    assert any(param.grad is not None for param in model.parameters())
+
+    param_ids = [id(param) for param in model.parameters()]
+    param_layouts = [_layout_signature(param) for param in model.parameters()]
+    buffer_layouts = [_layout_signature(buffer) for buffer in model.buffers()]
+    state_layouts = [_layout_signature(tensor) for tensor in _optimizer_tensors(optimizer)]
+    expected_params = [_local(param).detach().cpu().clone() for param in model.parameters()]
+    expected_buffers = [_local(buffer).detach().cpu().clone() for buffer in model.buffers()]
+    expected_states = [tensor.detach().cpu().clone() for tensor in _optimizer_tensors(optimizer)]
+    expected_shadow = [tensor.detach().cpu().clone() for tensor in ema.shadow]
+    fsdp_params = _fsdp_params(model)
+    storage_layouts = [_layout_signature(param._sharded_param_data) for param in fsdp_params]
+    expected_storage = [param._sharded_param_data.detach().cpu().clone() for param in fsdp_params]
+    has_padding = torch.tensor(
+        any(param._sharded_param_data.numel() > param.sharded_param.to_local().numel() for param in fsdp_params),
+        device=device,
+    )
+    dist.all_reduce(has_padding, op=dist.ReduceOp.MAX)
+    assert has_padding.item()
+    sleep_times = []
+    wake_times = []
+    awake_allocated = torch.cuda.memory_allocated()
+    sleep_allocated = awake_allocated
+
+    for _ in range(2):
+        start = time.perf_counter()
+        move_torch_model(model, "cpu", pin_memory=True)
+        move_torch_optimizer(optimizer, "cpu", pin_memory=True)
+        ema.to("cpu", non_blocking=True, pin_memory=True)
+        torch.cuda.synchronize()
+        sleep_times.append(time.perf_counter() - start)
+        sleep_allocated = torch.cuda.memory_allocated()
+        assert sleep_allocated < awake_allocated
+
+        assert [id(param) for param in model.parameters()] == param_ids
+        assert [_layout_signature(param) for param in model.parameters()] == param_layouts
+        assert [_layout_signature(buffer) for buffer in model.buffers()] == buffer_layouts
+        assert [_layout_signature(tensor) for tensor in _optimizer_tensors(optimizer)] == state_layouts
+        assert [_layout_signature(param._sharded_param_data) for param in fsdp_params] == storage_layouts
+        assert all(param.grad is None for param in model.parameters())
+        _assert_device(model.parameters(), "cpu", pinned=True)
+        _assert_device(model.buffers(), "cpu", pinned=True)
+        _assert_device(_optimizer_tensors(optimizer), "cpu", pinned=True)
+        _assert_device(ema.shadow, "cpu", pinned=True)
+        _assert_device((param._sharded_param_data for param in fsdp_params), "cpu", pinned=True)
+        assert model.state_dict()
+        ema.update()
+
+        start = time.perf_counter()
+        move_torch_model(model, device)
+        move_torch_optimizer(optimizer, device)
+        ema.to(device, non_blocking=True)
+        torch.cuda.synchronize()
+        wake_times.append(time.perf_counter() - start)
+
+        assert [id(param) for param in model.parameters()] == param_ids
+        assert [_layout_signature(param) for param in model.parameters()] == param_layouts
+        assert [_layout_signature(buffer) for buffer in model.buffers()] == buffer_layouts
+        assert [_layout_signature(tensor) for tensor in _optimizer_tensors(optimizer)] == state_layouts
+        assert [_layout_signature(param._sharded_param_data) for param in fsdp_params] == storage_layouts
+        _assert_device(model.parameters(), "cuda")
+        _assert_device(model.buffers(), "cuda")
+        _assert_device(_optimizer_tensors(optimizer), "cuda")
+        _assert_device(ema.shadow, "cuda")
+
+    for actual, expected in zip(model.parameters(), expected_params, strict=True):
+        torch.testing.assert_close(_local(actual).cpu(), expected)
+    for actual, expected in zip(model.buffers(), expected_buffers, strict=True):
+        torch.testing.assert_close(_local(actual).cpu(), expected)
+    for actual, expected in zip(_optimizer_tensors(optimizer), expected_states, strict=True):
+        torch.testing.assert_close(actual.cpu(), expected)
+    for actual, expected in zip(ema.shadow, expected_shadow, strict=True):
+        torch.testing.assert_close(actual.cpu(), expected)
+    for param, expected in zip(fsdp_params, expected_storage, strict=True):
+        torch.testing.assert_close(param._sharded_param_data.cpu(), expected)
+
+    optimizer.zero_grad(set_to_none=True)
+    model(inputs).sum().backward()
+    optimizer.step()
+    torch.cuda.synchronize()
+
+    dist.destroy_process_group()
+    if local_rank == 0:
+        print(
+            f"sleep_ms={sum(sleep_times) / len(sleep_times) * 1e3:.3f} "
+            f"wake_ms={sum(wake_times) / len(wake_times) * 1e3:.3f} "
+            f"awake_allocated={awake_allocated} sleep_allocated={sleep_allocated}"
+        )
+        print("OK")
+
+
+if __name__ == "__main__":
+    main()