chatlearn/models/fsdp_module.py

# Copyright 2024 Alibaba Group Holding Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""FSDP  module"""
import os
import random
import functools
import gc

import ray
import numpy as np
import torch
import torch.distributed as dist
from torch import optim, nn
from torch.distributed.fsdp import (MixedPrecision, ShardingStrategy, ShardedOptimStateDictConfig, ShardedStateDictConfig,
                                    FullStateDictConfig, StateDictType, FullyShardedDataParallel as FSDP)
from torch.distributed.fsdp.wrap import (size_based_auto_wrap_policy, transformer_auto_wrap_policy, lambda_auto_wrap_policy,
                                        _or_policy)
from torch.distributed.fsdp._runtime_utils import _lazy_init
from torch.multiprocessing.reductions import reduce_tensor

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.trainer_pt_utils import get_module_class_from_name

from chatlearn.utils.logger import debug_rank_0, log_rank_0
from chatlearn.utils.utils import dict_to_simplenamespace

from .torch_module import TorchModule


class FSDPModule(TorchModule):
    """TorchModule is the class for Alignment Torch models.

    Args
    ----
    name : str
        model name
    """

    def __init__(self, *args, **kwargs):

        super().__init__(*args, **kwargs)

        if not self.trainable:
            # inference only
            if (
                self.model_args.get("train_micro_batch_size")
                != self.module_args.generation_batch_size
            ):
                self._logger.info(
                    f"{self.name} Overwrite train_micro_batch_size with generation_batch_size {self.module_args.generation_batch_size}"
                )
            self.train_micro_batch_size = self.module_args.generation_batch_size
        else:
            self.train_micro_batch_size = self.runtime_args.train_micro_batch_size
            self.train_global_batch_size = self.runtime_args.train_global_batch_size

        self.fsdp_size = self.module_args.fsdp_size
        self.device_mesh = None

    def get_visible_gpus(self):
        """
        :meta private:
        """
        return ray.get_gpu_ids()

    @staticmethod
    def init_fn(x: torch.nn.Module):
        if torch.distributed.get_rank() != 0:
            x = x.to_empty(device=torch.cuda.current_device(), recurse=False)
            torch.cuda.empty_cache()
        return x

    @staticmethod
    def get_fsdp_wrap_policy(module:torch.nn.Module, min_num_params:int=0):
        """Get FSDP wrap policy for the module.

        Args:
            module: The module to get wrap policy for
            min_num_params: size based wrap policy min num params
        """

        default_transformer_cls_names_to_wrap = getattr(module, "_no_split_modules", None)
        fsdp_transformer_layer_cls_to_wrap = default_transformer_cls_names_to_wrap
        auto_wrap_policy = None

        policies = []

        # Add lambda policy for LoRA modules if is_lora is True
        # min_num_params must be 0 to use transformer_auto_wrap_policy
        if min_num_params > 0:
            size_policy = functools.partial(size_based_auto_wrap_policy, min_num_params=min_num_params)
            policies.append(size_policy)
        elif fsdp_transformer_layer_cls_to_wrap is not None:
            transformer_cls_to_wrap = set()
            for layer_class in fsdp_transformer_layer_cls_to_wrap:
                transformer_cls = get_module_class_from_name(module, layer_class)
                assert transformer_cls is not None, "Could not find the transformer layer class to wrap in the model."
                transformer_cls_to_wrap.add(transformer_cls)

            transformer_policy = functools.partial(
                transformer_auto_wrap_policy,
                transformer_layer_cls=transformer_cls_to_wrap,
            )
            policies.append(transformer_policy)
        ### hardcode for qwen2.5, fsdp warp for get submodule state dict ###
        def lambda_fn(sub_module: nn.Module):
            if sub_module in [module.model.embed_tokens, module.model.norm, module.lm_head]:
                return True
            return False
        lambda_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_fn)
        policies.append(lambda_policy)
        ## hardcode for qwen2.5, fsdp warp for get submodule state dict ##

        if len(policies) > 0:
            auto_wrap_policy = functools.partial(_or_policy, policies=policies)

        return auto_wrap_policy


    def create_device_mesh(self, world_size, fsdp_size):
        if not self.device_mesh:
            if world_size == fsdp_size:
                self.device_mesh = dist.device_mesh.init_device_mesh(
                    "cuda", mesh_shape=(world_size,), mesh_dim_names=["fsdp"]
                )
            else:
                self.device_mesh = dist.device_mesh.init_device_mesh(
                    "cuda",
                    mesh_shape=(fsdp_size, world_size // fsdp_size),
                    mesh_dim_names=["fsdp", "ddp"],
                )
            print(f"world size {world_size}, fsdp_size {fsdp_size}, {self.device_mesh}")

    def setup_distributed(self):
        print(self.get_dist_env())
        if not dist.is_initialized():
            dist.init_process_group(backend="nccl")

        self.create_device_mesh(self.world_size, self.fsdp_size)

    def peak_memory(self):
        """
        :meta private:
        """
        self._peak_memory = max(
            self._peak_memory, torch.cuda.max_memory_allocated() / (1024**3)
        )
        return self._peak_memory

    def empty_cache(self):
        """
        :meta private:
        """
        if not self.timers("empty_cache").started_:
            self.timers("empty_cache").start()
        peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
        debug_rank_0(
            f"{self.name} replica: {self.replica_id}, before empty cache, peak mem: {peak_mem:.2f} GiB",
            self._logger,
        )
        # Manual gc
        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
        debug_rank_0(
            f"{self.name} replica: {self.replica_id}, after empty cache, peak mem: {peak_mem:.2f} GiB",
            self._logger,
        )
        self.timers("empty_cache").stop()

    def check_param_exists(self, names):
        """
        check if the given names exists in current model

        :meta private:
        """
        not_exists = []
        for name in names:
            if not self.exist_parameter(name):
                not_exists.append(name)
        if not_exists:
            log_rank_0(
                f"parameters not exists: {not_exists} in model {self.name}",
                self._logger,
            )
            return False
        return True

    def create_model(self, model_path, torch_dtype):
        model = AutoModelForCausalLM.from_pretrained(
            pretrained_model_name_or_path=model_path,
            torch_dtype=torch_dtype,
            attn_implementation="flash_attention_2",
            trust_remote_code=True,
        )
        return model

    @property
    def data_parallel_size(self):
        """
        :meta private:
        """
        return dist.get_world_size()

    @property
    def data_parallel_rank(self):
        """
        :meta private:
        """
        return dist.get_rank()

    def tensor_parallel_rank(self):
        return self.data_parallel_rank

    def pipeline_parallel_rank(self):
        return 1

    def expert_model_parallel_size(self):
        return 1

    def model_setup(self):
        """
        :meta private:
        """
        super().model_setup()
        self.setup_distributed()
        args = dict_to_simplenamespace(self.model_args)
        self.args = args
        model = self.create_model(args.pretrain_or_model, torch_dtype=torch.bfloat16)
        self.tokenizer = AutoTokenizer.from_pretrained(
            args.pretrain_or_model, trust_remote_code=True, use_fast=True
        )
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False})
        mix_precision_config = MixedPrecision(
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.float32,
            buffer_dtype=torch.float32,
        )
        sharding_strategy = ShardingStrategy.FULL_SHARD
        auto_wrap_policy = self.get_fsdp_wrap_policy(module = model)
        self.model = FSDP(
            model,
            cpu_offload=None,
            auto_wrap_policy=auto_wrap_policy,
            device_id=torch.cuda.current_device(),
            sharding_strategy=sharding_strategy,
            mixed_precision=mix_precision_config,
            sync_module_states=True,
            param_init_fn=FSDPModule.init_fn,
            device_mesh=self.device_mesh,
            forward_prefetch=False,
        )
        self.model.to(torch.float32)
        FSDP.set_state_dict_type(self.model, StateDictType.SHARDED_STATE_DICT)
        if not self.trainable:
            self.optimizer = None
            self.model.eval()
        else:
            self.optimizer = optim.AdamW(
                self.model.parameters(),
                lr=self.module_args.args_dict.get("learning_rate", 2e-6),
                betas=(0.9, 0.999),
                weight_decay=1e-2
            )

        # resume model weights
        if self.resume_training:
            self.load_checkpoint(self._episode_id)

        self.offload()

    def get_fsdp_param_name(self):
        name_list = []
        for name, _ in self.model.named_parameters():
            parts = name.split('.')
            filtered_parts = [
                part for part in parts
                if part not in {"_fsdp_wrapped_module", "_flat_param"}
            ]
            cleaned_name = '.'.join(filtered_parts)
            name_list.append(cleaned_name)
        return name_list

    def get_weight_ipc_handles_by_name(self, block_name: str):
        """
        get fsdp warpped module weight by name get from named_parameters
        avoid get total model state_dict
        """
        torch.cuda.empty_cache()
        for prefix_name, module in self.model.named_modules():
            prefix_name = prefix_name.replace('_fsdp_wrapped_module.', '')
            if isinstance(module, FSDP) and prefix_name==block_name:
                state_dict = module.state_dict()
                reduce_tensor_dict = {}
                for name, param in state_dict.items():
                    reduce_tensor_dict['.'.join([prefix_name, name])] = reduce_tensor(param.full_tensor())
                return reduce_tensor_dict

    @torch.no_grad()
    def onload_weights(self,empty_cache=True):
        _lazy_init(self.model, self.model)
        assert self.model._is_root
        device_id = torch.cuda.current_device()
        for handle in self.model._all_handles:
            if handle._offload_params:
                continue
            flat_param = handle.flat_param
            handle.flat_param_to(torch.device(f"cuda:{device_id}"), non_blocking=True)
            # the following still keeps id(._local_shard) != id(.data)
            flat_param._local_shard = flat_param.data
        if empty_cache:
            torch.cuda.empty_cache()

    @torch.no_grad()
    def offload_weights(self, empty_cache=True):
        assert isinstance(self.model, FSDP)
        # lazy init FSDP model
        _lazy_init(self.model, self.model)
        assert self.model._is_root, "Only support root model offloading to CPU"
        for handle in self.model._all_handles:
            if handle._offload_params:
                continue
            flat_param = handle.flat_param
            assert (
                flat_param.data.data_ptr() == flat_param._local_shard.data_ptr()
                and id(flat_param.data) != id(flat_param._local_shard)
                and flat_param.data.size() == flat_param._local_shard.size()
            )
            handle.flat_param_to(torch.device("cpu"), non_blocking=True)
            # Explicit call to free unshard flat param
            handle._free_unsharded_flat_param()
            # the following still keeps id(._local_shard) != id(.data)
            flat_param._local_shard = flat_param.data
            assert id(flat_param._local_shard) != id(flat_param.data)

        # Explicit releas ipc handles
        torch.cuda.ipc_collect()
        if empty_cache:
            torch.cuda.empty_cache()

    @torch.no_grad()
    def offload_optimizer_states(self, empty_cache=True):
        if not self.optimizer.state:
            return
        for param_group in self.optimizer.param_groups:
            for param in param_group["params"]:
                state = self.optimizer.state[param]
                for key, value in state.items():
                    if isinstance(value, torch.Tensor):
                        state[key] = value.to("cpu", non_blocking=True)

        if empty_cache:
            torch.cuda.empty_cache()

    @torch.no_grad()
    def onload_optimizer_states(self, empty_cache=True):
        if not self.optimizer.state:
            return
        device = torch.cuda.current_device()
        for param_group in self.optimizer.param_groups:
            for param in param_group["params"]:
                state = self.optimizer.state[param]
                for key, value in state.items():
                    if isinstance(value, torch.Tensor):
                        state[key] = value.to(device, non_blocking=True)

        if empty_cache:
            torch.cuda.empty_cache()

    def save_checkpoint(self, iteration):
        save_dir = f"{self.runtime_args.output_dir}/save_model/{self.name}/{iteration}"
        if self.data_parallel_rank == 0 and not os.path.exists(save_dir):
            os.makedirs(save_dir)

        state_dict_cfg = ShardedStateDictConfig(offload_to_cpu=True)
        optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=True)
        with FSDP.state_dict_type(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):
            model_state_dict = self.model.state_dict()
            optimizer_state_dict = self.optimizer.state_dict() if self.optimizer is not None else None
            # lr_scheduler_state_dict = self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None
            extra_state_dict = {
                # "lr_scheduler": lr_scheduler_state_dict,
                "rng": self.get_rng_state(),
            }
            model_path = os.path.join(save_dir, f"model_world_size_{self.data_parallel_size}_rank_{self.data_parallel_rank}.pt")
            optim_path = os.path.join(save_dir, f"optim_world_size_{self.data_parallel_size}_rank_{self.data_parallel_rank}.pt")
            extra_path = os.path.join(save_dir, f"extra_state_world_size_{self.data_parallel_size}_rank_{self.data_parallel_rank}.pt")
            torch.save(model_state_dict, model_path)
            torch.save(optimizer_state_dict, optim_path)
            torch.save(extra_state_dict, extra_path)
        torch.distributed.barrier()

        # save for hf format
        if self.model_args.get("save_hf", True):
            state_dict_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
            with FSDP.state_dict_type(self.model, StateDictType.FULL_STATE_DICT, state_dict_cfg, None):
                model_state_dict = self.model.state_dict()
                if self.data_parallel_rank == 0:
                    hf_path = os.path.join(save_dir, "huggingface")
                    os.makedirs(hf_path, exist_ok=True)
                    model_config = self.model._fsdp_wrapped_module.config
                    model_config.save_pretrained(hf_path)
                    self.tokenizer.save_pretrained(hf_path)

                    with torch.device("meta"):
                        save_model = AutoModelForCausalLM.from_config(model_config, torch_dtype=torch.bfloat16)
                    save_model.to_empty(device="cpu")
                    save_model.save_pretrained(hf_path, state_dict=model_state_dict)
            self._logger.info(f"save checkpoint to {save_dir}")

    def load_checkpoint(self, iteration):
        load_dir = f"{self.runtime_args.output_dir}/save_model/{self.name}/{iteration}"
        if not os.path.exists(load_dir):
            self._logger.info(f"{load_dir} not exists, will skip load")
            return
        model_path = os.path.join(load_dir, f"model_world_size_{self.data_parallel_size}_rank_{self.data_parallel_rank}.pt")
        optim_path = os.path.join(load_dir, f"optim_world_size_{self.data_parallel_size}_rank_{self.data_parallel_rank}.pt")
        extra_state_path = os.path.join(load_dir, f"extra_state_world_size_{self.data_parallel_size}_rank_{self.data_parallel_rank}.pt")

        model_state_dict = torch.load(model_path, weights_only=False)
        optimizer_state_dict = torch.load(optim_path, weights_only=False)
        extra_state_dict = torch.load(extra_state_path, weights_only=False)

        state_dict_cfg = ShardedStateDictConfig(offload_to_cpu=True)
        optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=True)
        with FSDP.state_dict_type(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):
            self.model.load_state_dict(model_state_dict)
            if self.optimizer is not None:
                self.optimizer.load_state_dict(optimizer_state_dict)
        # recover random state
        if "rng" in extra_state_dict:
            # 'rng' may not exist for backward compatibility
            self.load_rng_state(extra_state_dict["rng"])
        torch.distributed.barrier()


    @staticmethod
    def get_rng_state():
        rng_state = {
            "cpu": torch.get_rng_state(),
            "cuda": torch.cuda.get_rng_state(),
            "numpy": np.random.get_state(),
            "random": random.getstate(),
        }
        return rng_state

    @staticmethod
    def load_rng_state(rng_state):
        torch.set_rng_state(rng_state["cpu"])
        torch.cuda.set_rng_state(rng_state["cuda"])
        np.random.set_state(rng_state["numpy"])
        random.setstate(rng_state["random"])