Spaces:
Running on Zero
Running on Zero
| """FSDP readers/writers and one scalar EMA; no implicit fabric overrides.""" | |
| import os | |
| import tempfile | |
| from datetime import timedelta | |
| from functools import partial | |
| import torch | |
| import torch.distributed as dist | |
| from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict | |
| from torch.distributed.fsdp import FullyShardedDataParallel as FSDP | |
| from torch.distributed.fsdp import MixedPrecision, ShardingStrategy | |
| from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy | |
| from src.utils.checkpoint_state import ( | |
| normalize_generator_fsdp_key, | |
| normalize_generator_fsdp_state_dict, | |
| ) | |
| def launch_distributed_job(): | |
| """Direct Python uses a local process group; torchrun uses its own rendezvous.""" | |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) | |
| torch.cuda.set_device(local_rank) | |
| if dist.is_initialized(): | |
| return | |
| kwargs = { | |
| "backend": "nccl", | |
| "device_id": torch.device("cuda", local_rank), | |
| "timeout": timedelta(minutes=30), | |
| } | |
| if "RANK" in os.environ: | |
| dist.init_process_group(**kwargs) | |
| else: | |
| descriptor, filename = tempfile.mkstemp(prefix="looking-ahead-pg-") | |
| os.close(descriptor) | |
| os.unlink(filename) | |
| dist.init_process_group(init_method=f"file://{filename}", rank=0, world_size=1, **kwargs) | |
| def barrier(): | |
| if dist.is_initialized(): | |
| dist.barrier() | |
| def is_main_process(): | |
| return not dist.is_initialized() or dist.get_rank() == 0 | |
| def fsdp_state_dict(model): | |
| """All ranks collect; only global rank zero retains ordinary weight tensors.""" | |
| checkpoint = get_model_state_dict( | |
| model, options=StateDictOptions(full_state_dict=True, cpu_offload=True) | |
| ) | |
| return checkpoint if is_main_process() else {} | |
| def fsdp_wrap(module, *, sharding_strategy, mixed_precision, sync_module_states=False): | |
| precision = ( | |
| MixedPrecision( | |
| param_dtype=torch.bfloat16, | |
| reduce_dtype=torch.float32, | |
| buffer_dtype=torch.float32, | |
| cast_forward_inputs=False, | |
| ) | |
| if mixed_precision | |
| else None | |
| ) | |
| strategy = { | |
| "full": ShardingStrategy.FULL_SHARD, | |
| "hybrid_full": ShardingStrategy.HYBRID_SHARD, | |
| "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2, | |
| "no_shard": ShardingStrategy.NO_SHARD, | |
| }[sharding_strategy] | |
| return FSDP( | |
| module, | |
| auto_wrap_policy=partial(size_based_auto_wrap_policy, min_num_params=int(5e7)), | |
| sharding_strategy=strategy, | |
| mixed_precision=precision, | |
| device_id=torch.cuda.current_device(), | |
| limit_all_gathers=True, | |
| use_orig_params=True, | |
| sync_module_states=sync_module_states, | |
| ) | |
| def _fsdp_sharding_strategy(module: torch.nn.Module): | |
| strategy = getattr(module, "sharding_strategy", None) | |
| if strategy is not None: | |
| return strategy | |
| fsdp_state = getattr(module, "_fsdp_state", None) | |
| if fsdp_state is not None: | |
| return getattr(fsdp_state, "sharding_strategy", None) | |
| return None | |
| def _effective_no_shard(module: torch.nn.Module) -> bool: | |
| """Return true when all FSDP wrappers use effective ``NO_SHARD``. | |
| A plain module is treated like ``NO_SHARD`` for EMA serialization helpers; | |
| this keeps the direct save/load path testable without requiring CUDA/FSDP. | |
| """ | |
| fsdp_modules = [submodule for submodule in module.modules() if isinstance(submodule, FSDP)] | |
| if not fsdp_modules: | |
| return True | |
| return all( | |
| _fsdp_sharding_strategy(submodule) == ShardingStrategy.NO_SHARD | |
| for submodule in fsdp_modules | |
| ) | |
| class EMA_FSDP: | |
| """ | |
| Exponential Moving Average for an FSDP-wrapped model. | |
| Design notes | |
| ------------ | |
| * **Sharded, on-GPU shadow.** Each rank stores an EMA shadow of *only its | |
| own local shard* of every ``flat_param`` (FSDP1) / ``DTensor`` (FSDP2), | |
| keeping the buffer in fp32 on the same device as the live parameter. | |
| This means ``update`` is a pure local elementwise op -- no | |
| ``summon_full_params``, no all-gather, no host-device traffic, no | |
| cross-rank synchronization -- so it is essentially free compared to | |
| the optimizer step. | |
| * **Collective only on save.** ``full_state_dict`` is the only path that | |
| materializes full EMA tensors for checkpointing. Sharded FSDP uses | |
| ``summon_full_params``; effective ``NO_SHARD`` returns the already-full | |
| local EMA shadows directly. It should only run at checkpoint time, never | |
| per-step. | |
| * Non-floating-point buffers (e.g. ``num_batches_tracked``) are never | |
| averaged. | |
| """ | |
| def __init__( | |
| self, | |
| fsdp_module: torch.nn.Module, | |
| decay: float = 0.999, | |
| ): | |
| self.decay = decay | |
| self.num_updates = 0 | |
| # name -> fp32 tensor on the *same device & with the same local | |
| # shape* as the corresponding live (sharded) parameter. Each rank | |
| # only holds its own shard. | |
| self.shadow: dict = {} | |
| self._is_main = is_main_process() | |
| self._init_shadow(fsdp_module) | |
| # ------------------------------------------------------------------ | |
| # Helpers | |
| # ------------------------------------------------------------------ | |
| def _iter_named_tensors(self, module): | |
| yield from module.named_parameters() | |
| # ------------------------------------------------------------------ | |
| # Initialization | |
| # ------------------------------------------------------------------ | |
| def _init_shadow(self, fsdp_module): | |
| # No collective: each rank just clones its own shard in fp32. For | |
| # FSDP1 the local shard is a 1-D ``flat_param`` slice; for FSDP2 it | |
| # is a ``DTensor`` whose ``_local_tensor`` we read transparently | |
| # (``detach().clone()`` on a DTensor preserves the DTensor wrapping, | |
| # which is exactly what we want for arithmetic later). | |
| for n, t in self._iter_named_tensors(fsdp_module): | |
| self.shadow[n] = t.detach().to(dtype=torch.float32).clone() | |
| # ------------------------------------------------------------------ | |
| # Update -- per-step hot path | |
| # ------------------------------------------------------------------ | |
| def update(self, fsdp_module): | |
| """Per-rank local EMA update. **No** collective communication. | |
| Cost is one fp32 mul_+add_ over the local shard, which is dominated | |
| by memory bandwidth and negligible next to the optimizer step. | |
| """ | |
| self.num_updates += 1 | |
| d = self.decay | |
| one_minus_d = 1.0 - d | |
| for n, t in self._iter_named_tensors(fsdp_module): | |
| shadow = self.shadow.get(n) | |
| if shadow is None: | |
| # New / unseen parameter -- initialize on the fly. | |
| self.shadow[n] = t.detach().to(dtype=torch.float32).clone() | |
| continue | |
| # Cast live param to fp32 *on device* (cheap if already fp32, | |
| # otherwise a fused kernel). No host transfer, no collective. | |
| src = t.detach() | |
| if src.dtype != torch.float32: | |
| src = src.to(dtype=torch.float32) | |
| shadow.mul_(d).add_(src, alpha=one_minus_d) | |
| # ------------------------------------------------------------------ | |
| # (De)serialization helpers | |
| # ------------------------------------------------------------------ | |
| def state_dict(self): | |
| """Return the *local-shard* shadow on every rank. | |
| Note: this is **not** a full state dict. Use :meth:`full_state_dict` | |
| if you want the rank-0-only fully-gathered version (e.g. for | |
| checkpointing). Each rank's tensors live on its own device. | |
| """ | |
| return self.shadow | |
| def load_full_state_dict( | |
| self, full_sd, fsdp_module, *, allow_missing_aux=False, expected_shapes=None | |
| ): | |
| """Restore ordinary EMA from each rank's validated full source state. | |
| All-rank writeback avoids PyTorch's unsupported rank0-only writeback. | |
| The temporary swap restores the original independent live shards. | |
| """ | |
| full_sd = normalize_generator_fsdp_state_dict(full_sd) | |
| expected = { | |
| normalize_generator_fsdp_key(name): tensor | |
| for name, tensor in self._iter_named_tensors(fsdp_module) | |
| } | |
| missing = set(expected) - set(full_sd) | |
| unexpected = set(full_sd) - set(expected) | |
| if unexpected or (missing and not allow_missing_aux): | |
| raise RuntimeError("EMA checkpoint has missing or unexpected parameter keys.") | |
| if allow_missing_aux: | |
| from src.utils.predictive_training_checkpoint import is_predictive_aux_key | |
| if any(not is_predictive_aux_key(k) for k in missing): | |
| raise RuntimeError("EMA checkpoint is missing backbone parameters") | |
| if expected_shapes is not None: | |
| for name, value in full_sd.items(): | |
| if tuple(value.shape) != expected_shapes[name]: | |
| raise RuntimeError(f"EMA checkpoint shape mismatch: {name}") | |
| if _effective_no_shard(fsdp_module): | |
| for name, tensor in expected.items(): | |
| if name in full_sd and full_sd[name].shape != tensor.shape: | |
| raise RuntimeError(f"EMA checkpoint shape mismatch: {name}") | |
| new_shadow = {} | |
| for n, t in self._iter_named_tensors(fsdp_module): | |
| canonical_name = normalize_generator_fsdp_key(n) | |
| src = full_sd.get(canonical_name, t) | |
| new_shadow[n] = ( | |
| src.detach() | |
| .to( | |
| device=t.device, | |
| dtype=torch.float32, | |
| ) | |
| .clone() | |
| ) | |
| self.shadow = new_shadow | |
| return | |
| # Backup live shards (each rank saves its own). | |
| backup = {n: t.data.clone() for n, t in self._iter_named_tensors(fsdp_module)} | |
| try: | |
| with FSDP.summon_full_params( | |
| fsdp_module, writeback=True, rank0_only=False, offload_to_cpu=False | |
| ): | |
| for name, tensor in self._iter_named_tensors(fsdp_module): | |
| canonical_name = normalize_generator_fsdp_key(name) | |
| if canonical_name in full_sd: | |
| tensor.copy_( | |
| full_sd[canonical_name].to(dtype=tensor.dtype, device=tensor.device) | |
| ) | |
| new_shadow = { | |
| name: tensor.detach().to(dtype=torch.float32).clone() | |
| for name, tensor in self._iter_named_tensors(fsdp_module) | |
| } | |
| self.shadow = new_shadow | |
| finally: | |
| for name, tensor in self._iter_named_tensors(fsdp_module): | |
| tensor.data.copy_(backup[name]) | |
| def full_state_dict(self, fsdp_module): | |
| """Return the *full* EMA state dict on rank 0 (empty dict elsewhere). | |
| This should only be used at checkpoint time. Effective ``NO_SHARD`` | |
| returns the already-full EMA shadows directly. Sharded FSDP temporarily | |
| swaps the live FSDP parameters for the EMA shadows, summons them in | |
| full on rank 0, copies them out, then restores the live weights -- all | |
| without persisting any change. | |
| """ | |
| if _effective_no_shard(fsdp_module): | |
| if not self._is_main: | |
| return {} | |
| return { | |
| normalize_generator_fsdp_key(n): shadow.detach() | |
| .to("cpu", dtype=torch.float32) | |
| .clone() | |
| for n, shadow in self.shadow.items() | |
| } | |
| # Clone the live local shards before copying EMA into their storage. | |
| # ``use_orig_params=True`` exposes views into FSDP's FlatParameter, so | |
| # retaining ``t.data`` here would only retain an alias. FSDP may reuse | |
| # and overwrite that storage while summoning full parameters, turning | |
| # an apparent restore into an accidental live <- EMA switch. | |
| backup = {} | |
| for n, t in self._iter_named_tensors(fsdp_module): | |
| shadow = self.shadow.get(n) | |
| if shadow is None: | |
| continue | |
| backup[n] = t.detach().clone() | |
| t.copy_(shadow.to(device=t.device, dtype=t.dtype), non_blocking=True) | |
| # CPU FSDP parameters are already host-resident. Asking PyTorch to | |
| # offload a rank0-only summon back to the same device can hang under | |
| # Gloo, while CUDA training still benefits from the original offload. | |
| offload_to_cpu = any(t.device.type != "cpu" for t in backup.values()) | |
| try: | |
| with FSDP.summon_full_params( | |
| fsdp_module, | |
| writeback=False, | |
| rank0_only=True, | |
| offload_to_cpu=offload_to_cpu, | |
| ): | |
| if not self._is_main: | |
| out = {} | |
| else: | |
| out = {} | |
| for n, t in self._iter_named_tensors(fsdp_module): | |
| out[normalize_generator_fsdp_key(n)] = ( | |
| t.detach().to(device="cpu", dtype=torch.float32).clone() | |
| ) | |
| finally: | |
| # Restore the independent live-shard copies regardless of what | |
| # happened above. Copying into the post-context parameter views | |
| # preserves FSDP's storage/view invariants. | |
| for n, t in self._iter_named_tensors(fsdp_module): | |
| if n in backup: | |
| t.copy_(backup[n], non_blocking=True) | |
| return out | |