chore: import upstream snapshot with attribution

2026-07-13 12:31:40 +08:00
commit 4df75a30cd
234 changed files with 41407 additions and 0 deletions
@@ -0,0 +1,149 @@
+from datetime import timedelta
+from functools import partial
+import os
+import torch
+import torch.distributed as dist
+from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType
+from torch.distributed.fsdp.api import CPUOffload
+from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
+
+
+def fsdp_state_dict(model):
+    fsdp_fullstate_save_policy = FullStateDictConfig(
+        offload_to_cpu=True, rank0_only=True
+    )
+    with FSDP.state_dict_type(
+        model, StateDictType.FULL_STATE_DICT, fsdp_fullstate_save_policy
+    ):
+        checkpoint = model.state_dict()
+
+    return checkpoint
+
+
+def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", min_num_params=int(5e7), transformer_module=None, cpu_offload=False):
+    if mixed_precision:
+        mixed_precision_policy = MixedPrecision(
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.float32,
+            buffer_dtype=torch.float32,
+            cast_forward_inputs=False
+        )
+    else:
+        mixed_precision_policy = None
+
+    if wrap_strategy == "transformer":
+        auto_wrap_policy = partial(
+            transformer_auto_wrap_policy,
+            transformer_layer_cls=transformer_module
+        )
+    elif wrap_strategy == "size":
+        auto_wrap_policy = partial(
+            size_based_auto_wrap_policy,
+            min_num_params=min_num_params
+        )
+    else:
+        raise ValueError(f"Invalid wrap strategy: {wrap_strategy}")
+
+    os.environ["NCCL_CROSS_NIC"] = "1"
+
+    sharding_strategy = {
+        "full": ShardingStrategy.FULL_SHARD,
+        "hybrid_full": ShardingStrategy.HYBRID_SHARD,
+        "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,
+        "no_shard": ShardingStrategy.NO_SHARD,
+    }[sharding_strategy]
+
+    module = FSDP(
+        module,
+        auto_wrap_policy=auto_wrap_policy,
+        sharding_strategy=sharding_strategy,
+        mixed_precision=mixed_precision_policy,
+        device_id=torch.cuda.current_device(),
+        limit_all_gathers=True,
+        use_orig_params=True,
+        cpu_offload=CPUOffload(offload_params=cpu_offload),
+        sync_module_states=False  # Load ckpt on rank 0 and sync to other ranks
+    )
+    return module
+
+
+def barrier():
+    if dist.is_initialized():
+        dist.barrier()
+
+
+def launch_distributed_job(backend: str = "nccl"):
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    host = os.environ["MASTER_ADDR"]
+    port = int(os.environ["MASTER_PORT"])
+
+    if ":" in host:  # IPv6
+        init_method = f"tcp://[{host}]:{port}"
+    else:  # IPv4
+        init_method = f"tcp://{host}:{port}"
+    # Use a long timeout so that slow collectives during checkpoint saving
+    # (e.g. FSDP.optim_state_dict all-gather + rank0-only disk write for a
+    # multi-GB full optimizer state) do not trip the NCCL watchdog on other
+    # ranks while they wait at the post-save barrier.
+    dist.init_process_group(rank=rank, world_size=world_size, backend=backend,
+                            init_method=init_method, timeout=timedelta(minutes=60))
+    torch.cuda.set_device(local_rank)
+
+
+class EMA_FSDP:
+    def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):
+        self.decay = decay
+        self.shadow = {}
+        self._init_shadow(fsdp_module)
+
+    @staticmethod
+    def _clean_param_name(name: str) -> str:
+        """Remove FSDP wrapper prefixes from parameter names."""
+        return name.replace("_fsdp_wrapped_module.", "").replace("_checkpoint_wrapped_module.", "").replace("_orig_mod.", "")
+
+    @torch.no_grad()
+    def _init_shadow(self, fsdp_module):
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=False):
+            for n, p in fsdp_module.module.named_parameters():
+                # Clean the parameter name to remove FSDP prefixes
+                # This ensures shadow keys are compatible with unwrapped models for inference
+                cleaned_name = self._clean_param_name(n)
+                self.shadow[cleaned_name] = p.detach().clone().float().cpu()
+
+    @torch.no_grad()
+    def update(self, fsdp_module):
+        d = self.decay
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=False):
+            for n, p in fsdp_module.module.named_parameters():
+                cleaned_name = self._clean_param_name(n)
+                if cleaned_name in self.shadow:
+                    self.shadow[cleaned_name].mul_(d).add_(p.detach().float().cpu(), alpha=1. - d)
+
+    # Optional helpers ---------------------------------------------------
+    def state_dict(self):
+        # Return shadow dict directly - keys are already cleaned during init/update
+        # This makes the state_dict directly usable for inference with unwrapped models
+        return self.shadow            # picklable
+
+    def load_state_dict(self, sd):
+        # Handle both cases: with or without FSDP prefixes
+        # This ensures backward compatibility and flexibility
+        cleaned_sd = {}
+        for k, v in sd.items():
+            # Remove FSDP prefixes if present to match internal naming convention
+            cleaned_key = self._clean_param_name(k)
+            cleaned_sd[cleaned_key] = v.clone()
+        self.shadow = cleaned_sd
+
+    def copy_to(self, fsdp_module):
+        # load EMA weights into an (unwrapped) copy of the generator
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=True):
+            for n, p in fsdp_module.module.named_parameters():
+                cleaned_name = self._clean_param_name(n)
+                if cleaned_name in self.shadow:
+                    p.data.copy_(self.shadow[cleaned_name].to(p.dtype, device=p.device))