ALSv/self-forcing

1from datetime import timedelta2from functools import partial3import os4import torch5import torch.distributed as dist6from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType7from torch.distributed.fsdp.api import CPUOffload8from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy9 10 11def fsdp_state_dict(model):12    fsdp_fullstate_save_policy = FullStateDictConfig(13        offload_to_cpu=True, rank0_only=True14    )15    with FSDP.state_dict_type(16        model, StateDictType.FULL_STATE_DICT, fsdp_fullstate_save_policy17    ):18        checkpoint = model.state_dict()19 20    return checkpoint21 22 23def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", min_num_params=int(5e7), transformer_module=None, cpu_offload=False):24    if mixed_precision:25        mixed_precision_policy = MixedPrecision(26            param_dtype=torch.bfloat16,27            reduce_dtype=torch.float32,28            buffer_dtype=torch.float32,29            cast_forward_inputs=False30        )31    else:32        mixed_precision_policy = None33 34    if wrap_strategy == "transformer":35        auto_wrap_policy = partial(36            transformer_auto_wrap_policy,37            transformer_layer_cls=transformer_module38        )39    elif wrap_strategy == "size":40        auto_wrap_policy = partial(41            size_based_auto_wrap_policy,42            min_num_params=min_num_params43        )44    else:45        raise ValueError(f"Invalid wrap strategy: {wrap_strategy}")46 47    os.environ["NCCL_CROSS_NIC"] = "1"48 49    sharding_strategy = {50        "full": ShardingStrategy.FULL_SHARD,51        "hybrid_full": ShardingStrategy.HYBRID_SHARD,52        "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,53        "no_shard": ShardingStrategy.NO_SHARD,54    }[sharding_strategy]55 56    module = FSDP(57        module,58        auto_wrap_policy=auto_wrap_policy,59        sharding_strategy=sharding_strategy,60        mixed_precision=mixed_precision_policy,61        device_id=torch.cuda.current_device(),62        limit_all_gathers=True,63        use_orig_params=True,64        cpu_offload=CPUOffload(offload_params=cpu_offload),65        sync_module_states=False  # Load ckpt on rank 0 and sync to other ranks66    )67    return module68 69 70def barrier():71    if dist.is_initialized():72        dist.barrier()73 74 75def launch_distributed_job(backend: str = "nccl"):76    rank = int(os.environ["RANK"])77    local_rank = int(os.environ["LOCAL_RANK"])78    world_size = int(os.environ["WORLD_SIZE"])79    host = os.environ["MASTER_ADDR"]80    port = int(os.environ["MASTER_PORT"])81 82    if ":" in host:  # IPv683        init_method = f"tcp://[{host}]:{port}"84    else:  # IPv485        init_method = f"tcp://{host}:{port}"86    dist.init_process_group(rank=rank, world_size=world_size, backend=backend,87                            init_method=init_method, timeout=timedelta(minutes=30))88    torch.cuda.set_device(local_rank)89 90 91class EMA_FSDP:92    def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):93        self.decay = decay94        self.shadow = {}95        self._init_shadow(fsdp_module)96 97    @torch.no_grad()98    def _init_shadow(self, fsdp_module):99        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP100        with FSDP.summon_full_params(fsdp_module, writeback=False):101            for n, p in fsdp_module.module.named_parameters():102                self.shadow[n] = p.detach().clone().float().cpu()103 104    @torch.no_grad()105    def update(self, fsdp_module):106        d = self.decay107        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP108        with FSDP.summon_full_params(fsdp_module, writeback=False):109            for n, p in fsdp_module.module.named_parameters():110                self.shadow[n].mul_(d).add_(p.detach().float().cpu(), alpha=1. - d)111 112    # Optional helpers ---------------------------------------------------113    def state_dict(self):114        return self.shadow            # picklable115 116    def load_state_dict(self, sd):117        self.shadow = {k: v.clone() for k, v in sd.items()}118 119    def copy_to(self, fsdp_module):120        # load EMA weights into an (unwrapped) copy of the generator121        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP122        with FSDP.summon_full_params(fsdp_module, writeback=True):123            for n, p in fsdp_module.module.named_parameters():124                if n in self.shadow:125                    p.data.copy_(self.shadow[n].to(p.dtype, device=p.device))126