Zhangzhe's Blog

import math
import torch
from torch.optim import Optimizer
import torch.distributed as dist
from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
class AdaLomo(Optimizer):
    """
    一个自定义的优化器类AdaLomo，用于在分布式训练中的梯度更新。
    该类实现两个梯度更新函数 :meth:`fuse_update` 和 :meth:`fuse_update_zero3`，分别用于非ZeRO和ZeRO模式下的梯度更新。
    :param model: 待优化的模型
    :param lr: 学习率，默认值为1e-3
    :param eps: 正则化系数。eps[0]防止梯度平方太小，eps[1]用于在根据参数的RMS放缩学习率时防止步长太大
    :param clip_threshold: 归一化update矩阵时的阈值
    :param decay_rate: 梯度平方移动平均的衰减率
    :param clip_grad_norm: 梯度裁剪的范数阈值
        .. note::
            clip_grad_norm须为正数
    :param clip_grad_value: 梯度裁剪的值域阈值
    :param weight_decay: 权重衰减系数，默认值为0.0
    :param loss_scale: 损失缩放系数，可以用来提高训练精度，但是太大可能会导致nan
    """
    def __init__(
        self,
        model,
        lr=1e-3,
        loss_scale=2 ** 10,
        eps=(1e-30, 1e-3),
        clip_threshold=1.0,
        decay_rate=-0.8,
        clip_grad_norm=None,
        clip_grad_value=None,
        weight_decay=0.0,
    ):
        self.model = model
        self.lr = lr
        self.clip_grad_norm = clip_grad_norm
        self.clip_grad_value = clip_grad_value
        self.weight_decay = weight_decay
        self.loss_scale = loss_scale
        if self.weight_decay > 0.0:
            self.do_weight_decay = True
        else:
            self.do_weight_decay = False
        self.eps = eps
        self.step_num = 0
        self.decay_rate = decay_rate
        self.clip_threshold = clip_threshold
        # for grad norm
        if self.clip_grad_norm is not None and self.clip_grad_norm <= 0:
            raise ValueError(
                f"clip_grad_norm should be positive, got {self.clip_grad_norm}."
            )
        self.gather_norm = False
        self.grad_norms = []
        self.clip_coef = None
        # check if zero3 is enabled
        self.zero3_enabled = is_deepspeed_zero3_enabled()
        if self.zero3_enabled:  # zero3 is enabled
            self.grad_func = self.fuse_update_zero3()
        else:
            self.grad_func = self.fuse_update()
        self.exp_avg_sq = {}
        self.exp_avg_sq_row = {}
        self.exp_avg_sq_col = {}
        # register hook function, which will be called through the backward process
        for n, p in self.model.named_parameters():
            if self.zero3_enabled:
                if len(p.ds_shape) == 1:
                    self.exp_avg_sq[n] = torch.zeros(
                        p.ds_shape[0], dtype=torch.float32
                    ).cuda()
                else:
                    self.exp_avg_sq_row[n] = torch.zeros(
                        p.ds_shape[0], dtype=torch.float32
                    ).cuda()
                    self.exp_avg_sq_col[n] = torch.zeros(
                        p.ds_shape[1], dtype=torch.float32
                    ).cuda()
            else:
                if len(p.data.shape) == 1:
                    self.exp_avg_sq[n] = torch.zeros(
                        p.data.shape[0], dtype=torch.float32
                    ).cuda()
                else:
                    self.exp_avg_sq_row[n] = torch.zeros(
                        p.data.shape[0], dtype=torch.float32
                    ).cuda()
                    self.exp_avg_sq_col[n] = torch.zeros(
                        p.data.shape[1], dtype=torch.float32
                    ).cuda()
            if p.requires_grad:
                p.register_hook(self.grad_func)
        defaults = dict(
            lr=lr,
            eps=eps,
            weight_decay=weight_decay,
            clip_grad_norm=clip_grad_norm,
            clip_grad_value=clip_grad_value,
        )
        super(AdaLomo, self).__init__(self.model.parameters(), defaults)
    @staticmethod
    def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
        # copy from fairseq's adafactor implementation:
        # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505
        r_factor = (
            (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True))
            .rsqrt_()
            .unsqueeze(-1)
        )
        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
        return torch.mul(r_factor, c_factor)
    @staticmethod
    def _rms(tensor):
        return tensor.norm(2) / (tensor.numel() ** 0.5)
    def fuse_update(self):
        """
        在非ZeRO模式下更新模型参数的梯度。
        :return: func，一个闭包函数，用于更新模型参数的梯度
        """
        def func(x):
            """
            闭包函数，用于更新模型参数的梯度。
            """
            with torch.no_grad():
                for n, p in self.model.named_parameters():
                    if p.requires_grad and p.grad is not None:
                        grad_fp32 = p.grad.to(torch.float32)
                        p.grad = None
                        if self.loss_scale:
                            grad_fp32.div_(self.loss_scale)
                        if self.gather_norm:
                            # we adopt two backward pass for gradient norm computation and parameter update, respectively.
                            self.grad_norms.append(torch.norm(grad_fp32, 2.0))
                        else:
                            # grad clip or norm
                            if (
                                self.clip_grad_value is not None
                                and self.clip_grad_value > 0
                            ):
                                # Clipping gradients by their value
                                grad_fp32.clamp_(
                                    min=-self.clip_grad_value, max=self.clip_grad_value
                                )
                            if (
                                self.clip_grad_norm is not None
                                and self.clip_grad_norm > 0
                                and self.clip_coef is not None
                            ):
                                # Normalize the gradient according to its norm (computed in another pass)
                                grad_fp32.mul_(self.clip_coef)
                            beta2t = 1.0 - math.pow(self.step_num, self.decay_rate)
                            update = (grad_fp32 ** 2) + self.eps[0]
                            if len(p.data.shape) > 1:
                                self.exp_avg_sq_row[n].mul_(beta2t).add_(
                                    update.mean(dim=-1), alpha=1.0 - beta2t
                                )
                                self.exp_avg_sq_col[n].mul_(beta2t).add_(
                                    update.mean(dim=-2), alpha=1.0 - beta2t
                                )
                                update = self._approx_sq_grad(
                                    self.exp_avg_sq_row[n], self.exp_avg_sq_col[n]
                                )
                                update.mul_(grad_fp32)
                            else:
                                self.exp_avg_sq[n].mul_(beta2t).add_(
                                    update, alpha=1.0 - beta2t
                                )
                                update = self.exp_avg_sq[n].rsqrt().mul_(grad_fp32)
                            update.div_(
                                (self._rms(update) / self.clip_threshold).clamp_(
                                    min=1.0
                                )
                            )
                            p_fp32 = p.data.to(torch.float32)
                            p_rms = torch.norm(p_fp32, 2.0) / math.sqrt(p.numel())
                            lr = self.lr
                            param_scale = max(self.eps[1], p_rms)
                            lr = lr * param_scale
                            if self.do_weight_decay:
                                p_fp32.mul_(1.0 - lr * self.weight_decay)
                            p_fp32.add_(update, alpha=-lr)
                            p.data.copy_(p_fp32)
            return x
        return func
    def fuse_update_zero3(self):
        """
        在ZeRO模式下更新模型参数的梯度。
        :return: func，一个闭包函数，用于更新模型参数的梯度。
        """
        def func(x):
            with torch.no_grad():
                for n, p in self.model.named_parameters():
                    if p.grad is not None:
                        torch.distributed.all_reduce(
                            p.grad, op=torch.distributed.ReduceOp.AVG, async_op=False
                        )
                        grad_fp32 = p.grad.to(torch.float32)
                        p.grad = None
                        if self.loss_scale:
                            grad_fp32.div_(self.loss_scale)
                        if self.gather_norm:
                            # we adopt two backward pass for gradient norm computation and parameter update, respectively.
                            self.grad_norms.append(torch.norm(grad_fp32, 2.0))
                        else:  # update param
                            partition_size = p.ds_tensor.numel()
                            start = partition_size * self.dp_rank
                            end = min(start + partition_size, grad_fp32.numel())
                            if self.clip_grad_value is not None:
                                # Clipping gradients by their value
                                grad_fp32.clamp_(
                                    min=-self.clip_grad_value, max=self.clip_grad_value
                                )
                            if (
                                self.clip_grad_norm is not None
                                and self.clip_grad_norm > 0
                                and self.clip_coef is not None
                            ):
                                # Normalize the gradient according to its norm (computed in another pass)
                                grad_fp32.mul_(self.clip_coef)
                            beta2t = 1.0 - math.pow(self.step_num, self.decay_rate)
                            update = (grad_fp32 ** 2) + self.eps[0]  # 改成addcmul_
                            if len(p.ds_shape) > 1:
                                self.exp_avg_sq_row[n].mul_(beta2t).add_(
                                    update.mean(dim=-1), alpha=1.0 - beta2t
                                )
                                self.exp_avg_sq_col[n].mul_(beta2t).add_(
                                    update.mean(dim=-2), alpha=1.0 - beta2t
                                )
                                update = self._approx_sq_grad(
                                    self.exp_avg_sq_row[n], self.exp_avg_sq_col[n]
                                )
                                update.mul_(grad_fp32)
                            else:
                                self.exp_avg_sq[n].mul_(beta2t).add_(
                                    update, alpha=1.0 - beta2t
                                )
                                update = self.exp_avg_sq[n].rsqrt().mul_(grad_fp32)
                            update.div_(
                                (self._rms(update) / self.clip_threshold).clamp_(
                                    min=1.0
                                )
                            )
                            one_dim_update = update.view(-1)
                            partitioned_update = one_dim_update.narrow(
                                0, start, end - start
                            )
                            param_fp32 = p.ds_tensor.to(torch.float32)
                            partitioned_p = param_fp32.narrow(0, 0, end - start)
                            p_rms = torch.norm(partitioned_p, 2.0) ** 2
                            dist.all_reduce(p_rms, op=torch.distributed.ReduceOp.SUM)
                            p_rms = (p_rms / p.ds_numel).sqrt()
                            lr = self.lr
                            param_scale = max(self.eps[1], p_rms)
                            lr = lr * param_scale
                            if self.do_weight_decay:
                                partitioned_p.mul_(1.0 - lr * self.weight_decay)
                            partitioned_p.add_(partitioned_update, alpha=-lr)
                            p.ds_tensor[: end - start] = partitioned_p
            return x
        return func
    def fused_backward(self, loss, lr):
        """
        执行一步反向传播并更新模型的梯度。
        :param loss: 模型的loss值
        :param lr: 学习率
        """
        self.lr = lr
        if self.loss_scale:
            loss = loss * self.loss_scale
        self.step_num += 1
        loss.backward()
        # update the last parameter since the last parameter in the computaiton graph is not ready when calling hook functions
        # the argument of grad_func is just a placeholder, and it can be anything.
        self.grad_func(0)
    def grad_norm(self, loss):
        """
        计算梯度的范数。
        :param loss: 模型的loss值
        """
        self.gather_norm = True
        self.grad_norms = []
        if self.loss_scale:
            loss = loss * self.loss_scale
        loss.backward(retain_graph=True)
        # update the last parameter since the last parameter in the computaiton graph is not ready when calling hook functions
        # the argument of grad_func is just a placeholder, and it can be anything.
        self.grad_func(0)
        with torch.no_grad():
            # The norm is computed over all gradients together, as if they were
            # concatenated into a single vector. Gradients are modified in-place.
            self.grad_norms = torch.stack(self.grad_norms)
            total_norm = torch.norm(self.grad_norms, 2.0)
            self.clip_coef = float(self.clip_grad_norm) / (total_norm + 1e-6)
            self.clip_coef = torch.clamp(self.clip_coef, max=1.0)
        self.gather_norm = False

class LOMO(Optimizer):
    """
    一个自定义的优化器类LOMO，用于在分布式训练中的梯度更新。
    该类实现两个梯度更新函数 :meth:`fuse_update` 和 :meth:`fuse_update_zero3`，分别用于非ZeRO和ZeRO模式下的梯度更新。
    :param model: 待优化的模型
    :param lr: 学习率，默认值为1e-3
    :param clip_grad_norm: 梯度裁剪的范数阈值
        .. note::
            clip_grad_norm须为正数
    :param clip_grad_value: 梯度裁剪的值域阈值
    """
    def __init__(self, model, lr=1e-3, clip_grad_norm=None, clip_grad_value=None):
        self.model = model
        self.lr = lr
        self.local_rank = int(os.environ["LOCAL_RANK"])
        self.world_size = dist.get_world_size()
        self.clip_grad_norm = clip_grad_norm
        self.clip_grad_value = clip_grad_value
        # for grad norm
        if self.clip_grad_norm is not None and self.clip_grad_norm <= 0:
            raise ValueError(
                f"clip_grad_norm should be positive, got {self.clip_grad_norm}."
            )
        self.gather_norm = False
        self.grad_norms = []
        self.clip_coef = None
        # check if zero3 is enabled
        p0 = list(self.model.parameters())[0]
        if hasattr(p0, "ds_tensor"):  # zero3 is enabled
            self.grad_func = self.fuse_update_zero3()
        else:
            self.grad_func = self.fuse_update()
        # check if fp16 is enabled
        if p0.dtype == torch.float16:
            self.loss_scaler = DynamicLossScaler(
                init_scale=2 ** 16,
            )  # TODO: add args
            if self.clip_grad_norm is None:
                raise ValueError(
                    "Loss scaling is recommended to be used with grad norm to get better performance."
                )
        else:
            self.loss_scaler = None
        # register hook function, which will be called through the backward process
        for n, p in self.model.named_parameters():
            if p.requires_grad:
                p.register_hook(self.grad_func)
        defaults = dict(
            lr=lr, clip_grad_norm=clip_grad_norm, clip_grad_value=clip_grad_value
        )
        super(LOMO, self).__init__(self.model.parameters(), defaults)
    def fuse_update(self):
        """
        在非ZeRO模式下更新模型参数的梯度。
        :return: func，一个闭包函数，用于更新模型参数的梯度
        """
        def func(x):
            """
            闭包函数，用于更新模型参数的梯度。
            """
            with torch.no_grad():
                for n, p in self.model.named_parameters():
                    if p.requires_grad and p.grad is not None:
                        if self.loss_scaler:
                            if (
                                self.loss_scaler.has_overflow_serial
                                or self.loss_scaler._has_inf_or_nan(p.grad)
                            ):
                                # if the overflow is detected, drop the gradient
                                p.grad = None
                                self.loss_scaler.has_overflow_serial = True
                                break
                        grad_fp32 = p.grad.to(torch.float32)
                        p.grad = None
                        if self.loss_scaler:
                            grad_fp32.div_(self.loss_scaler.loss_scale)
                        if self.gather_norm:
                            # we adopt two backward pass for gradient norm compuation and parameter update, respectively.
                            self.grad_norms.append(torch.norm(grad_fp32, 2.0))
                        else:
                            if (
                                self.clip_grad_value is not None
                                and self.clip_grad_value > 0
                            ):
                                # Clipping gradients by their value
                                grad_fp32.clamp_(
                                    min=-self.clip_grad_value, max=self.clip_grad_value
                                )
                            if (
                                self.clip_grad_norm is not None
                                and self.clip_grad_norm > 0
                                and self.clip_coef is not None
                            ):
                                # Normalize the gradient according to its norm (computed in another pass)
                                grad_fp32.mul_(self.clip_coef)
                            p_fp32 = p.data.to(torch.float32)
                            p_fp32.add_(grad_fp32, alpha=-self.lr)
                            p.data.copy_(p_fp32)
            return x
        return func
    def fuse_update_zero3(self):
        """
        在ZeRO模式下更新模型参数的梯度。
        :return: func，一个闭包函数，用于更新模型参数的梯度。
        """
        def func(x):
            with torch.no_grad():
                for n, p in self.model.named_parameters():
                    if p.grad is not None:
                        torch.distributed.all_reduce(
                            p.grad, op=torch.distributed.ReduceOp.AVG, async_op=False
                        )
                        if self.loss_scaler:
                            if (
                                self.loss_scaler.has_overflow_serial
                                or self.loss_scaler._has_inf_or_nan(p.grad)
                            ):
                                # if the overflow is detected, drop the gradient
                                p.grad = None
                                self.loss_scaler.has_overflow_serial = True
                                break
                        grad_fp32 = p.grad.to(torch.float32)
                        p.grad = None
                        param_fp32 = p.ds_tensor.to(torch.float32)
                        if self.loss_scaler:
                            grad_fp32.div_(self.loss_scaler.loss_scale)
                        if self.gather_norm:
                            # we adopt two backward pass for gradient norm compuation and parameter update, respectively.
                            self.grad_norms.append(torch.norm(grad_fp32, 2.0))
                        else:  # update param
                            one_dim_grad_fp32 = grad_fp32.view(-1)
                            partition_size = p.ds_tensor.numel()
                            start = partition_size * self.local_rank
                            end = min(start + partition_size, grad_fp32.numel())
                            partitioned_grad_fp32 = one_dim_grad_fp32.narrow(
                                0, start, end - start
                            )
                            if self.clip_grad_value is not None:
                                # Clipping gradients by their value
                                partitioned_grad_fp32.clamp_(
                                    min=-self.clip_grad_value, max=self.clip_grad_value
                                )
                            if (
                                self.clip_grad_norm is not None
                                and self.clip_grad_norm > 0
                                and self.clip_coef is not None
                            ):
                                # Normalize the gradient according to its norm (computed in another pass)
                                partitioned_grad_fp32.mul_(self.clip_coef)
                            partitioned_p = param_fp32.narrow(0, 0, end - start)
                            partitioned_p.add_(partitioned_grad_fp32, alpha=-self.lr)
                            p.ds_tensor[: end - start] = partitioned_p
            return x
        return func
    def fused_backward(self, loss, lr):
        """
        执行一步反向传播并更新模型的梯度（真正计算梯度和更新参数）。
        :param loss: 模型的loss值
        :param lr: 学习率
        """
        self.lr = lr
        # Users need call grad_norm themselves and then call backward_step
        if (
            self.clip_grad_norm is not None
            and self.clip_grad_norm > 0
            and self.clip_coef is None
        ):
            raise ValueError(
                "clip_grad_norm is not None, but clip_coef is None. "
                "Please call optimizer.grad_norm() before optimizer.fused_backward()."
            )
        if self.loss_scaler:
            loss = loss * self.loss_scaler.loss_scale
        loss.backward()
        # update the last parameter since the last parameter in the computaiton graph is not ready when calling hook functions
        # the argument of grad_func is just a placeholder, and it can be anything.
        self.grad_func(0)
    def grad_norm(self, loss):
        """
        计算梯度的范数（虽然做了一次 forward + backward，但实际上只是用于计算梯度的范数，后续做 clip grad norm 用）。
        :param loss: 模型的loss值
        """
        self.gather_norm = True
        self.grad_norms = []
        if self.loss_scaler:
            self.loss_scaler.has_overflow_serial = False
            loss = loss * self.loss_scaler.loss_scale
        loss.backward(retain_graph=True)
        # update the last parameter since the last parameter in the computaiton graph is not ready when calling hook functions
        # the argument of grad_func is just a placeholder, and it can be anything.
        self.grad_func(0)
        if self.loss_scaler and self.loss_scaler.has_overflow_serial:
            self.loss_scaler.update_scale(overflow=True)
            with torch.no_grad():  # clear gradients
                for n, p in self.model.named_parameters():
                    p.grad = None
            return
        with torch.no_grad():
            # The norm is computed over all gradients together, as if they were
            # concatenated into a single vector. Gradients are modified in-place.
            self.grad_norms = torch.stack(self.grad_norms)
            total_norm = torch.norm(self.grad_norms, 2.0)
            self.clip_coef = float(self.clip_grad_norm) / (total_norm + 1e-6)
            self.clip_coef = torch.clamp(self.clip_coef, max=1.0)
        self.gather_norm = False

import os
import torch
import torch.distributed as dist
from datasets import load_dataset
from datasets import Split
from transformers import AutoTokenizer, AutoModelForCausalLM
from torch.utils.data import IterableDataset, DataLoader
import deepspeed
from torch.utils.tensorboard import SummaryWriter
# 初始化 DeepSpeed
deepspeed.init_distributed()
# 设置环境变量以避免 tokenizers 并行化问题
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# 设置分布式训练所需的环境变量
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
os.environ["RANK"] = str(local_rank)
os.environ["WORLD_SIZE"] = str(world_size)
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "23333"
# NCCL 超时参数配置
os.environ["NCCL_BLOCKING_WAIT"] = "1"
os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "1"
os.environ["NCCL_DEBUG"] = "INFO"
os.environ["NCCL_TIMEOUT"] = "600"
# 初始化进程组
if not dist.is_initialized():
    dist.init_process_group(backend="nccl")
# config
model_path = "./models/meta-llama/Llama-3.2-3B"
dataset_path = "./dataset/LLM/yahma___alpaca-cleaned"
ds_config = "./ds_config.json"
log_dir = "./logs"
max_length = 1024
batch_size = 8
train_epochs = 3
# Load the local alpaca dataset
dataset_stream = load_dataset(dataset_path, split=Split.TRAIN, streaming=True)
# tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path, token=True)
tokenizer.pad_token = tokenizer.eos_token  # llama-3.2 tokenizer padding token not set
# model
model = AutoModelForCausalLM.from_pretrained(model_path, token=True)
model.to(f"cuda:{local_rank}")  # 将模型移动到对应的 GPU
# dataset
class StreamDataset(IterableDataset):
    def __init__(self, data_stream, tokenizer, max_length=max_length):
        self.data_stream = data_stream
        self.tokenizer = tokenizer
        self.max_length = max_length
    def __iter__(self):
        for data in self.data_stream:
            instruction = data["instruction"]
            input_text = data["input"]
            output = data["output"]
            # 构造 Prompt
            if input_text.strip():
                prompt = f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n"
            else:
                prompt = f"### Instruction:\n{instruction}\n\n### Response:\n"
            # 拼接完整输入序列：prompt + output
            full_text = prompt + output
            # Tokenize
            input_ids = tokenizer(
                full_text,
                truncation=True,
                max_length=self.max_length,
                padding="max_length",
            )["input_ids"]
            # 构造 labels: 仅监督 output 部分
            # 找到 output 起始位置
            prompt_length = len(
                tokenizer(prompt, padding="do_not_pad", truncation=False)["input_ids"]
            )
            labels = [-100] * prompt_length + input_ids[prompt_length:]  # 忽略 prompt 的部分
            labels = labels[: self.max_length]  # 截断
            labels = labels + [-100] * (max_length - len(labels))  # Padding 保持一致长度
            yield {
                "input_ids": torch.tensor(input_ids, dtype=torch.long),
                "labels": torch.tensor(labels, dtype=torch.long),
            }
    def __len__(self):
        return 51760  # cleaned alpaca dataset rows
def train():
    # DataLoader
    train_dataset = StreamDataset(dataset_stream, tokenizer)
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size)
    # TensorBoard SummaryWriter
    writer = SummaryWriter(log_dir=log_dir) if dist.get_rank() == 0 else None
    # DeepSpeed 初始化
    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model,
        model_parameters=model.parameters(),
        config=ds_config,
        training_data=train_dataset,
    )
    # 训练循环
    for epoch in range(train_epochs):  # num_train_epochs
        for step, batch in enumerate(train_dataloader):
            inputs = {
                key: value.to(f"cuda:{local_rank}") for key, value in batch.items()
            }
            outputs = model_engine(**inputs)  # 添加 labels 参数
            loss = outputs.loss
            model_engine.backward(loss)
            model_engine.step()
            # 日志打印和 TensorBoard 记录
            if dist.get_rank() == 0:
                if step % 10 == 0:
                    print(
                        f"Epoch: {epoch}, Step: {step}/{len(train_dataloader)}, Loss: {loss.item()}"
                    )
                    writer.add_scalar(
                        "Loss/train", loss.item(), epoch * len(train_dataloader) + step
                    )
        checkpoint_path = f"./checkpoints"
        model_engine.save_checkpoint(checkpoint_path, tag=f"epoch_{epoch}")
        print(f"Checkpoint saved at {checkpoint_path}")
    # 关闭 TensorBoard SummaryWriter
    if writer:
        writer.close()
if __name__ == "__main__":
    train()

import os
from transformers import AutoModelForCausalLM, AutoTokenizer
final_model_save_path = "./checkpoints/final_transformers_model"  # 保存最终模型的路径
original_model_path = "./models/meta-llama/Llama-3.2-3B"  # 原始模型的路径
sharded_model_dir = "./checkpoints/fp32_checkpoints"  # SFT fp32 分片模型的路径
# 1. 将原始模型配置复制到最终模型目录，需要让 transformers 自动识别模型
os.system(
    f"cp {os.path.join(original_model_path, 'config.json')} {final_model_save_path}"
)
# 2. 将分片模型转换为 transformers 模型
model = AutoModelForCausalLM.from_pretrained(sharded_model_dir)
model.save_pretrained(final_model_save_path)
# 3. 补全 tokenizer 配置
tokenizer = AutoTokenizer.from_pretrained(original_model_path, token=True)
tokenizer.save_pretrained(final_model_save_path)

import torch
from transformers import pipeline
model_id = "./checkpoints/final_transformers_model"
pipe = pipeline(
    "text-generation", model=model_id, torch_dtype=torch.bfloat16, device_map="auto"
)
instruction = "Use Python code to implement."
input_text = "Find prime numbers within 100."
total_input = f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n"
print(
    pipe(total_input, max_length=1024)[0][
        "generated_text"
    ]
)

class ActorCriticNetwork(nn.Module):
    """Actor-Critic 网络，包含策略网络和价值网络。"""
    def __init__(self, obs_space_dims: int, action_space_dims: int):
        """初始化 Actor-Critic 网络。
        参数：
            obs_space_dims: 观测空间维度
            action_space_dims: 动作空间维度
        """
        super().__init__()
        hidden_space1 = 16
        hidden_space2 = 32
        # 共享网络
        self.shared_net = nn.Sequential(
            nn.Linear(obs_space_dims, hidden_space1),
            nn.Tanh(),
            nn.Linear(hidden_space1, hidden_space2),
            nn.Tanh(),
        )
        # 策略网络输出动作的均值和标准差
        self.policy_mean_net = nn.Linear(hidden_space2, action_space_dims)
        self.policy_stddev_net = nn.Linear(hidden_space2, action_space_dims)
        # 价值网络输出状态价值估计
        self.value_net = nn.Linear(hidden_space2, 1)
    def forward(self, x: torch.Tensor):
        """前向传播，输出动作参数和状态价值。
        参数：
            x: 环境的观测值
        返回：
            action_means: 动作均值
            action_stddevs: 动作标准差
            state_values: 状态价值估计
        """
        shared_features = self.shared_net(x.float())
        action_means = self.policy_mean_net(shared_features)
        action_stddevs = torch.log(
            1 + torch.exp(self.policy_stddev_net(shared_features))
        )
        state_values = self.value_net(shared_features)
        return action_means, action_stddevs, state_values
class ActorCriticAgent:
    """Actor-Critic 算法的实现。"""
    def __init__(self, obs_space_dims: int, action_space_dims: int):
        """初始化 Agent。
        参数：
            obs_space_dims: 观测空间维度
            action_space_dims: 动作空间维度
        """
        self.learning_rate = 1e-4
        self.gamma = 0.99
        self.eps = 1e-6
        self.log_probs = []
        self.state_values = []
        self.rewards = []
        self.net = ActorCriticNetwork(obs_space_dims, action_space_dims)
        self.optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.learning_rate)
    def sample_action(self, state: np.ndarray):
        """根据策略采样动作。
        参数：
            state: 环境的观测值
        返回：
            action: 采样的动作
        """
        state = torch.tensor(np.array([state]))
        action_means, action_stddevs, state_value = self.net(state)
        distrib = Normal(action_means[0] + self.eps, action_stddevs[0] + self.eps)
        action = distrib.sample()
        log_prob = distrib.log_prob(action)
        action = action.numpy()
        self.log_probs.append(log_prob)
        self.state_values.append(state_value)
        return action
    def update(self):
        """更新策略网络和价值网络。"""
        # 将列表转换为张量
        log_probs = torch.stack(self.log_probs)
        state_values = torch.cat(self.state_values).squeeze()
        rewards = self.rewards
        # 计算折扣回报
        returns = []
        G = 0
        for R in rewards[::-1]:
            G = R + self.gamma * G
            returns.insert(0, G)
        returns = torch.tensor(returns)
        # 计算损失
        policy_loss = -(log_probs * returns.detach()).sum()
        value_loss = nn.functional.mse_loss(state_values, returns)
        loss = policy_loss + value_loss
        # 更新网络参数
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        # 清空存储的值
        self.log_probs = []
        self.state_values = []
        self.rewards = []

# Actor-Critic 算法的实现
actor_loss = -(log_probs * returns.detach()).sum()
# Advantage Actor-Critic 算法的实现
advantages = returns - state_values
actor_loss = -(log_probs * advantages.detach()).sum()