PyTorch 深度学习框架

import torch

# 从列表创建
a = torch.tensor([1.0, 2.0, 3.0])  # 1D 张量
b = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)  # 2D 张量，指定类型

# 常用工厂函数
zeros = torch.zeros(3, 4)        # 全零张量，形状 (3, 4)
ones = torch.ones(2, 3, 5)       # 全一张量
empty = torch.empty(2, 3)        # 未初始化张量（值不确定）
rand = torch.rand(2, 3)          # 均匀分布 [0, 1) 随机张量
randn = torch.randn(2, 3)        # 标准正态分布 N(0,1) 随机张量
arange = torch.arange(0, 10, 2)  # 等差序列：[0, 2, 4, 6, 8]
linspace = torch.linspace(0, 1, 5)  # 等间隔序列：[0, 0.25, 0.5, 0.75, 1.0]

# LLM 中常用的初始化方法
# Xavier/Glorot 初始化（适用于 Sigmoid/Tanh）
xavier = torch.empty(512, 512)
torch.nn.init.xavier_uniform_(xavier)

# Kaiming/He 初始化（适用于 ReLU，Transformer 中常用）
kaiming = torch.empty(512, 512)
torch.nn.init.kaiming_normal_(kaiming, mode='fan_in', nonlinearity='relu')

# 从 NumPy 转换
import numpy as np
np_array = np.array([1, 2, 3])
tensor_from_np = torch.from_numpy(np_array)  # 共享内存
np_from_tensor = a.numpy()  # CPU 张量转 NumPy，共享内存

x = torch.randn(2, 3, 4)

# 基本属性
x.shape       # torch.Size([2, 3, 4]) — 形状
x.dtype       # torch.float32 — 数据类型
x.device      # cpu 或 cuda:0 — 所在设备
x.ndim        # 3 — 维度数
x.numel()     # 24 — 元素总数

# 形状变换
x.view(2, 12)       # 返回新视图（共享内存），要求张量连续
x.reshape(2, 12)    # 优先返回视图，必要时拷贝
x.permute(2, 0, 1)  # 维度重排：(2,3,4) → (4,2,3)
x.transpose(0, 1)   # 交换两个维度
x.unsqueeze(0)      # 在维度0插入：(2,3,4) → (1,2,3,4)
x.squeeze()         # 移除所有大小为1的维度

# LLM 中的典型用法：注意力掩码维度扩展
# attention_mask: (batch_size, seq_len) → (batch_size, 1, 1, seq_len)
attention_mask = torch.ones(4, 128)  # batch=4, seq_len=128
expanded_mask = attention_mask[:, None, None, :]  # (4, 1, 1, 128)

x = torch.arange(12).reshape(3, 4)
# tensor([[ 0,  1,  2,  3],
#         [ 4,  5,  6,  7],
#         [ 8,  9, 10, 11]])

x[0]          # 第一行：[0, 1, 2, 3]
x[:, 1]       # 第二列：[1, 5, 9]
x[0:2, 1:3]   # 子矩阵：[[1,2],[5,6]]

# 高级索引
indices = torch.tensor([0, 2])
x[:, indices]     # 选取第0和第2列
x[x > 5]          # 布尔索引：[6, 7, 8, 9, 10, 11]

# LLM 常用：选取特定 token 的嵌入
# hidden_states: (batch, seq_len, hidden_dim)
hidden_states = torch.randn(2, 10, 768)
# 获取每个序列最后一个 token 的表示
last_token = hidden_states[:, -1, :]  # (2, 768)

# (3, 1) + (1, 4) → (3, 4)
a = torch.ones(3, 1)
b = torch.ones(1, 4)
c = a + b  # 形状 (3, 4)

# LLM 典型场景：RoPE 位置编码
# freqs: (seq_len, head_dim/2)  →  扩展为 (1, seq_len, 1, head_dim/2)
# query: (batch, seq_len, num_heads, head_dim)
freqs = torch.randn(128, 64)
freqs = freqs[None, :, None, :]  # (1, 128, 1, 64) — 自动广播到 batch 和 num_heads 维度

# requires_grad=True 标记需要计算梯度的张量
x = torch.tensor([2.0, 3.0], requires_grad=True)

# 前向计算（自动构建计算图）
y = x ** 2        # y = [4, 9]
z = y.sum()       # z = 13

# 反向传播
z.backward()

# 查看梯度
print(x.grad)     # dz/dx = 2x = [4, 6]

# 梯度累积：默认情况下梯度会累加
x = torch.tensor([2.0], requires_grad=True)

for _ in range(3):
    y = x ** 2
    y.backward()
    print(x.grad)  # 4.0 → 8.0 → 12.0（梯度累积！）

# 正确做法：每次反向传播前清零梯度
x = torch.tensor([2.0], requires_grad=True)
for _ in range(3):
    y = x ** 2
    if x.grad is not None:
        x.grad.zero_()  # 清零梯度
    y.backward()
    print(x.grad)  # 4.0 → 4.0 → 4.0（正确）

x = torch.tensor([2.0], requires_grad=True)

# no_grad：推理时禁用梯度计算，节省内存
with torch.no_grad():
    y = x ** 2  # y.requires_grad = False

# detach：从计算图中分离张量
y = x ** 2
z = y.detach()  # z 不再跟踪梯度，但与 y 共享数据

# torch.enable_grad：在 no_grad 内部临时启用梯度
with torch.no_grad():
    with torch.enable_grad():
        y = x ** 2  # 此处可以计算梯度

# 典型训练循环中的梯度管理
model = ...  # nn.Module
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

for batch in dataloader:
    # 1. 前向传播
    outputs = model(batch['input_ids'])
    loss = loss_fn(outputs, batch['labels'])

    # 2. 反向传播
    optimizer.zero_grad()    # 清零梯度（等价于 for p in model.parameters(): p.grad=None）
    loss.backward()          # 计算梯度

    # 3. 梯度裁剪（防止梯度爆炸，LLM 训练中常用）
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    # 4. 更新参数
    optimizer.step()

import torch.nn as nn

class SimpleMLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()  # 必须调用父类初始化
        # 在 __init__ 中定义所有层和参数
        self.fc1 = nn.Linear(input_dim, hidden_dim)   # 全连接层
        self.activation = nn.GELU()                     # 激活函数
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(p=0.1)               # Dropout 正则化

    def forward(self, x):
        # 在 forward 中定义前向计算逻辑
        x = self.fc1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 实例化并使用
model = SimpleMLP(input_dim=768, hidden_dim=3072, output_dim=768)
output = model(torch.randn(2, 10, 768))  # (batch, seq_len, output_dim)

# 访问所有参数
for name, param in model.named_parameters():
    print(f"{name}: shape={param.shape}, requires_grad={param.requires_grad}")

# 访问特定层的参数
print(model.fc1.weight.shape)  # (hidden_dim, input_dim) = (3072, 768)
print(model.fc1.bias.shape)    # (hidden_dim,) = (3072,)

# 冻结参数（LLM 微调中常用：冻结预训练层，只训练新层）
for name, param in model.named_parameters():
    if 'fc1' in name:  # 只冻结 fc1 层
        param.requires_grad = False

# 统计参数量
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"总参数量: {total_params:,}, 可训练参数量: {trainable_params:,}")

class MultiHeadSelfAttention(nn.Module):
    """简化版多头自注意力，展示 nn.Module 的典型用法"""
    def __init__(self, hidden_dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads
        assert hidden_dim % num_heads == 0, "hidden_dim 必须能被 num_heads 整除"

        self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.o_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)

    def forward(self, x, attention_mask=None):
        batch_size, seq_len, _ = x.shape

        # 线性投影并分头
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled Dot-Product Attention
        scale = self.head_dim ** 0.5
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / scale

        if attention_mask is not None:
            attn_weights = attn_weights.masked_fill(attention_mask == 0, float('-inf'))

        attn_weights = F.softmax(attn_weights, dim=-1)
        attn_output = torch.matmul(attn_weights, v)

        # 合并头
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        return self.o_proj(attn_output)

# 随机梯度下降（带动量和权重衰减）
optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.01,           # 学习率
    momentum=0.9,      # 动量系数，加速收敛
    weight_decay=1e-4  # L2 正则化（权重衰减）
)

# Adam 优化器（LLM 中最常用的优化器之一）
optimizer = torch.optim.Adam(
    model.parameters(),
    lr=1e-4,           # 学习率
    betas=(0.9, 0.999),  # (β1, β2) 一阶和二阶动量衰减系数
    eps=1e-8,          # 防止除零的小常数
    weight_decay=0.0   # 权重衰减
)

# AdamW：解耦权重衰减的 Adam（LLM 训练的标准选择）
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.95),   # LLM 中 β2 常用 0.95 而非 0.999
    eps=1e-8,
    weight_decay=0.01,   # 权重衰减系数
    amsgrad=False        # 是否使用 AMSGrad 变体
)

from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR

# 方式1：余弦退火（LLM 训练中最常用）
scheduler = CosineAnnealingLR(
    optimizer,
    T_max=100000,    # 周期长度（通常等于总训练步数）
    eta_min=1e-5     # 最小学习率
)

# 方式2：线性预热 + 余弦退火（LLM 标准配置）
warmup_scheduler = LinearLR(optimizer, start_factor=0.01, total_iters=2000)
cosine_scheduler = CosineAnnealingLR(optimizer, T_max=98000, eta_min=1e-5)
scheduler = SequentialLR(
    optimizer,
    schedulers=[warmup_scheduler, cosine_scheduler],
    milestones=[2000]  # 在第 2000 步切换调度器
)

# 训练循环中使用
for step, batch in enumerate(dataloader):
    loss = model(batch).loss
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    scheduler.step()  # 每个 step 后更新学习率

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_ddp():
    dist.init_process_group(backend="nccl")  # NCCL 后端用于 GPU 通信
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

def main():
    setup_ddp()
    local_rank = int(os.environ["LOCAL_RANK"])

    model = SimpleMLP(768, 3072, 768).cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for batch in dataloader:
        batch = {k: v.cuda(local_rank) for k, v in batch.items()}
        output = model(batch['input_ids'])
        loss = output.loss
        optimizer.zero_grad()
        loss.backward()        # DDP 自动同步梯度
        optimizer.step()

    dist.destroy_process_group()

# 启动命令：torchrun --nproc_per_node=4 train.py

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

model = SimpleMLP(768, 3072, 768).cuda(local_rank)

# FULL_SHARD：完全分片（最省显存，通信量最大）
# SHARD_GRAD_OP：仅分片梯度和优化器状态（折中方案）
# NO_SHARD：等同于 DDP
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    device_id=local_rank
)

from torch.cuda.amp import autocast, GradScaler

model = SimpleMLP(768, 3072, 768).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()  # 梯度缩放器，防止 FP16 梯度下溢

for batch in dataloader:
    batch = {k: v.cuda() for k, v in batch.items()}
    optimizer.zero_grad()

    # autocast 上下文管理器：自动选择 FP16/BF16 进行计算
    with autocast(dtype=torch.bfloat16):  # 或 torch.float16
        output = model(batch['input_ids'])
        loss = output.loss

    # 缩放损失 → 反向传播 → 反缩放梯度 → 更新参数
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)  # 反缩放，以便正确裁剪梯度
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)      # 如果梯度中含 inf/nan 则跳过此步
    scaler.update()             # 更新缩放因子

# BF16 (bfloat16)：指数位8位，尾数位7位 — 与 FP32 相同的数值范围
# FP16 (float16)：指数位5位，尾数位10位 — 更高精度但数值范围小

# BF16 不需要 GradScaler（不易溢出）
with autocast(dtype=torch.bfloat16):
    output = model(inputs)
    loss = output.loss
loss.backward()  # 直接反向传播，无需 scaler
optimizer.step()

# FP16 需要 GradScaler（容易溢出/下溢）
# 如上 3.6.1 示例所示

# 查看可用 GPU
print(torch.cuda.device_count())    # GPU 数量
print(torch.cuda.current_device())  # 当前设备 ID

# 张量设备迁移
x = torch.randn(2, 3)
x_gpu = x.to('cuda')             # CPU → GPU
x_gpu = x.cuda()                 # 等价写法
x_cpu = x_gpu.to('cpu')          # GPU → CPU
x_gpu1 = x_gpu.to('cuda:1')     # GPU 0 → GPU 1

# 模型设备迁移
model = SimpleMLP(768, 3072, 768)
model = model.to('cuda')

# 查看显存使用情况
print(torch.cuda.memory_allocated() / 1e9, "GB")      # 当前张量占用的显存
print(torch.cuda.memory_reserved() / 1e9, "GB")       # CUDA 缓存池占用的显存
print(torch.cuda.max_memory_allocated() / 1e9, "GB")  # 峰值显存

# 清理显存缓存
torch.cuda.empty_cache()  # 释放未使用的缓存（不影响仍在使用的张量）

# 手动垃圾回收
import gc
del large_tensor
gc.collect()
torch.cuda.empty_cache()

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class GPT2Block(nn.Module):
    def __init__(self, hidden_dim, num_heads, ffn_dim, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(hidden_dim)
        self.attn = nn.MultiheadAttention(hidden_dim, num_heads, dropout=dropout, batch_first=True)
        self.ln2 = nn.LayerNorm(hidden_dim)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_dim, ffn_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, hidden_dim),
            nn.Dropout(dropout),
        )

    def forward(self, x, attention_mask=None):
        # Pre-LN Transformer
        residual = x
        x = self.ln1(x)
        x, _ = self.attn(x, x, x, attn_mask=attention_mask, need_weights=False)
        x = residual + x

        residual = x
        x = self.ln2(x)
        x = self.ffn(x)
        x = residual + x
        return x

class GPT2Model(nn.Module):
    def __init__(self, vocab_size=50257, hidden_dim=768, num_layers=12,
                 num_heads=12, ffn_dim=3072, max_seq_len=1024):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, hidden_dim)
        self.position_embedding = nn.Embedding(max_seq_len, hidden_dim)
        self.blocks = nn.ModuleList([
            GPT2Block(hidden_dim, num_heads, ffn_dim) for _ in range(num_layers)
        ])
        self.ln_f = nn.LayerNorm(hidden_dim)
        self.lm_head = nn.Linear(hidden_dim, vocab_size, bias=False)

        # 权重共享
        self.lm_head.weight = self.token_embedding.weight

        # 因果掩码
        self.register_buffer("causal_mask",
            torch.triu(torch.ones(max_seq_len, max_seq_len), diagonal=1).bool())

    def forward(self, input_ids):
        batch_size, seq_len = input_ids.shape
        positions = torch.arange(seq_len, device=input_ids.device)

        x = self.token_embedding(input_ids) + self.position_embedding(positions)
        attn_mask = self.causal_mask[:seq_len, :seq_len]

        for block in self.blocks:
            x = block(x, attention_mask=attn_mask)

        x = self.ln_f(x)
        logits = self.lm_head(x)
        return logits

# 使用示例
model = GPT2Model()
input_ids = torch.randint(0, 50257, (2, 128))  # batch=2, seq_len=128
logits = model(input_ids)
print(logits.shape)  # (2, 128, 50257)

class LoRALinear(nn.Module):
    """LoRA：低秩适配器，冻结原始权重，仅训练低秩矩阵"""
    def __init__(self, original_linear, r=8, alpha=16):
        super().__init__()
        self.original = original_linear
        self.original.weight.requires_grad = False  # 冻结原始权重
        if self.original.bias is not None:
            self.original.bias.requires_grad = False

        in_dim = original_linear.in_features
        out_dim = original_linear.out_features

        self.lora_A = nn.Parameter(torch.randn(in_dim, r) * 0.01)
        self.lora_B = nn.Parameter(torch.zeros(r, out_dim))
        self.scaling = alpha / r

    def forward(self, x):
        # W'x = Wx + (B @ A) * scaling * x
        return self.original(x) + (x @ self.lora_A @ self.lora_B) * self.scaling

# 应用 LoRA 到模型
model = GPT2Model()
for name, module in model.named_modules():
    if isinstance(module, nn.Linear) and 'q_proj' in name:
        # 替换 Q 投影为 LoRA 版本
        parent_name, child_name = name.rsplit('.', 1)
        parent = model.get_submodule(parent_name)
        setattr(parent, child_name, LoRALinear(module, r=8, alpha=16))

# 1. 使用梯度检查点（用计算换显存）
from torch.utils.checkpoint import checkpoint

class GPT2BlockWithCheckpoint(nn.Module):
    def forward(self, x):
        # 不保存中间激活，反向传播时重新计算
        return checkpoint(self._forward, x)

    def _forward(self, x):
        # 原始前向逻辑
        ...

# 2. 使用 mixed precision 减少显存占用
with autocast(dtype=torch.bfloat16):
    output = model(inputs)  # 激活值和梯度用 BF16 存储，节省约 50% 显存

# 3. 及时释放不需要的张量
del output, loss
torch.cuda.empty_cache()

# 1. 梯度裁剪（LLM 训练必备）
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 2. 学习率预热（防止训练初期梯度爆炸）
scheduler = LinearLR(optimizer, start_factor=0.01, total_iters=2000)

# 3. 使用 BF16 代替 FP16（避免精度问题）
with autocast(dtype=torch.bfloat16):  # 推荐，A100/H100 原生支持
    ...

# 4. 检测 NaN/Inf
if torch.isnan(loss) or torch.isinf(loss):
    print("警告：检测到异常损失值，跳过此步")
    continue

from torch.utils.data import DataLoader, DistributedSampler

# 分布式训练中必须使用 DistributedSampler
sampler = DistributedSampler(dataset, shuffle=True)
dataloader = DataLoader(
    dataset,
    batch_size=8,            # 每卡 batch size
    sampler=sampler,         # 分布式采样器
    num_workers=4,           # 数据加载线程数
    pin_memory=True,         # 锁页内存，加速 CPU→GPU 传输
    drop_last=True,          # 丢弃最后不完整的 batch
)

# 每个 epoch 需要设置 sampler 的 epoch
for epoch in range(num_epochs):
    sampler.set_epoch(epoch)
    for batch in dataloader:
        ...

# 保存完整模型（不推荐，不灵活）
torch.save(model, 'model.pth')

# 推荐方式：只保存 state_dict
torch.save(model.state_dict(), 'model_state_dict.pth')

# 加载
model = GPT2Model()
model.load_state_dict(torch.load('model_state_dict.pth'))
model.eval()  # 切换到评估模式

# 分布式训练中保存（只由 rank 0 保存）
if dist.get_rank() == 0:
    torch.save(model.module.state_dict(), 'model.pth')