diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 332cfb09..dcf5db3b 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/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: diff --git a/miles/backends/fsdp_utils/ema.py b/miles/backends/fsdp_utils/ema.py index 9120ef3b..5a8161a1 100644 --- a/miles/backends/fsdp_utils/ema.py +++ b/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.""" diff --git a/tests/fast-gpu/backends/fsdp_utils/_sleep_wakeup_worker.py b/tests/fast-gpu/backends/fsdp_utils/_sleep_wakeup_worker.py new file mode 100644 index 00000000..029be211 --- /dev/null +++ b/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() diff --git a/tests/fast-gpu/backends/fsdp_utils/test_sleep_wakeup.py b/tests/fast-gpu/backends/fsdp_utils/test_sleep_wakeup.py new file mode 100644 index 00000000..887711a1 --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/test_sleep_wakeup.py @@ -0,0 +1,38 @@ +"""FSDP2 sleep/wakeup through pinned host memory.""" + +from tests.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=30, + suite="stage-b-5-gpu-h200", + labels=["fsdp"], +) + +import os +import subprocess +import sys +from pathlib import Path + +_WORKER = Path(__file__).with_name("_sleep_wakeup_worker.py") + + +def test_fsdp_sleep_wakeup_uses_pinned_memory(): + env = os.environ.copy() + env["PYTHONUNBUFFERED"] = "1" + result = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + "--nproc_per_node=2", + str(_WORKER), + ], + env=env, + capture_output=True, + text=True, + timeout=300, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "OK" in result.stdout