PyTorch 面试题

# Tensor 元数据
tensor = torch.tensor([[1, 2, 3],
                       [4, 5, 6]])
# .storage()  → 底层一维连续内存 [1, 2, 3, 4, 5, 6]
# .shape     → torch.Size([2, 3])
# .stride()  → (3, 1)   # 沿 dim0 步长 3，沿 dim1 步长 1
# .is_contiguous() → True

# Channels Last (NHWC) — 适合 CNN + cuDNN 加速
x_nhwc = x.contiguous(memory_format=torch.channels_last)
# 等价于 stride = (C, 1, W, C)  # NHWC 排列

# 场景 1：推理评估，节省显存
with torch.no_grad():
    y = model(x)

# 场景 2：从图中分离某中间变量
z = some_tensor.detach()  # 返回同数据的新 Tensor，requires_grad=False

# 场景 3：冻结参数
for p in model.parameters():
    p.requires_grad = False  # 参与前向但不参与梯度计算

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 1. nn.Parameter — 参与梯度更新的可学习参数
        self.weight = nn.Parameter(torch.randn(3, 3))

        # 2. nn.Module — 子模块，自动注册到 _modules
        self.linear = nn.Linear(3, 3)

        # 3. register_buffer — 不参与梯度但需随模型移动（device/.to()）
        self.register_buffer('running_mean', torch.zeros(3))

        # 4. 注册到 _parameters / _buffers / _modules 三个 OrderedDict

def init_weights(module):
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, nonlinearity='relu')
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.LayerNorm):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)

model.apply(init_weights)  # 递归应用到所有子模块

class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        # 单样本 IO 与预处理
        enc = self.tokenizer(self.texts[idx], truncation=True, max_length=512)
        return {
            'input_ids': torch.tensor(enc['input_ids']),
            'label': torch.tensor(self.labels[idx], dtype=torch.long)
        }

def collate_fn(batch):
    # 变长序列 padding
    input_ids = [item['input_ids'] for item in batch]
    padded = nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=0)
    labels = torch.stack([item['label'] for item in batch])
    return {'input_ids': padded, 'labels': labels}

class StreamDataset(IterableDataset):
    def __iter__(self):
        # 适用：超大文件、数据库流、网络流、TFRecord 多 shard
        worker_info = torch.utils.data.get_worker_info()
        for shard in self.shards[worker_info.id::worker_info.num_workers]:
            for sample in self._read_shard(shard):
                yield sample

# torchrun --nproc_per_node=4 train.py
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
model = MyModel().cuda(local_rank)
model = DDP(model, device_ids=[local_rank], find_unused_parameters=False)
sampler = DistributedSampler(dataset, shuffle=True)
loader = DataLoader(dataset, batch_size=32, sampler=sampler, drop_last=True)

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy

# auto_wrap_policy: 按层大小自动分片
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock}  # 自定义 Block 类
)

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3 等价
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    ),
    auto_wrap_policy=wrap_policy,
    backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
    device_id=local_rank,
)

# deepspeed --num_gpus=4 train.py
import deepspeed
model_engine, optimizer, _, scheduler = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config='ds_config.json'
)
# ds_config.json
{
  "train_batch_size": 32,
  "fp16": {"enabled": true},
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {"device": "cpu"},
    "offload_param": {"device": "cpu"}
  },
  "activation_checkpointing": {"partition_activations": true}
}

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()  # FP16 专用，BF16 无需 scaler

for x, y in loader:
    optimizer.zero_grad(set_to_none=True)
    with autocast(dtype=torch.float16):  # 自动算子级别精度选择
        y_pred = model(x)
        loss = criterion(y_pred, y)

    scaler.scale(loss).backward()  # loss × scale
    scaler.unscale_(optimizer)     # 梯度 ÷ scale
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 裁剪未缩放梯度
    scaler.step(optimizer)         # 检测 inf/nan
    scaler.update()                # 动态调整 scale

# BF16 无需 Loss Scaling，指数范围与 FP32 相同
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    y_pred = model(x)
    loss = criterion(y_pred, y)
loss.backward()
optimizer.step()

# === 保存 ===
checkpoint = {
    'epoch': epoch,
    'global_step': global_step,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'scaler_state_dict': scaler.state_dict() if scaler else None,
    'best_metric': best_metric,
    'rng_state': torch.get_rng_state(),  # CPU 随机状态
    'cuda_rng_state': torch.cuda.get_rng_state_all(),  # 所有 GPU RNG
}
torch.save(checkpoint, f'ckpt_epoch{epoch}.pt')

# === 加载（断点续训） ===
checkpoint = torch.load('ckpt_epoch10.pt', map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
if scaler and checkpoint['scaler_state_dict']:
    scaler.load_state_dict(checkpoint['scaler_state_dict'])
start_epoch = checkpoint['epoch'] + 1
global_step = checkpoint['global_step']
torch.set_rng_state(checkpoint['rng_state'])
torch.cuda.set_rng_state_all(checkpoint['cuda_rng_state'])

# DDP 模型的 state_dict 多 'module.' 前缀
# 方案 1：先剥前缀再 load
state_dict = {k.replace('module.', ''): v for k, v in ckpt.items()}
model.load_state_dict(state_dict)

# 方案 2：使用 model.module 访问
model.module.load_state_dict(ckpt)

# FSDP 加载（PyTorch 2.x 支持直接 load 完整状态）
from torch.distributed.fsdp import FullyShardedDataParallel, StateDictType
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
    model.load_state_dict(full_state_dict)  # rank 0 广播

# 1. 仅在 rank 0 保存
if dist.get_rank() == 0:
    torch.save(checkpoint, path)

# 2. 滚动保存最近 N 个
ckpts = sorted(glob('ckpt_*.pt'), key=os.path.getmtime)
for old in ckpts[:-3]:
    os.remove(old)

# 3. 保存 EMA 模型（推理质量更稳）
ema_model = torch.optim.swa_utils.AveragedModel(model)
torch.save(ema_model.module.state_dict(), 'ema.pt')

# 4. 异步保存避免 IO 阻塞
import threading
def async_save(obj, path):
    threading.Thread(target=torch.save, args=(obj, path), daemon=True).start()

torch.onnx.export(
    model,                          # nn.Module
    (dummy_input,),                 # 示例输入（tuple）
    "model.onnx",
    input_names=['input_ids'],      # ONNX 图输入名
    output_names=['logits'],
    dynamic_axes={                  # 动态维度
        'input_ids': {0: 'batch', 1: 'sequence'},
        'logits': {0: 'batch', 1: 'sequence'},
    },
    opset_version=17,               # ONNX 算子集版本
    do_constant_folding=True,       # 常量折叠优化
    training=torch.onnx.TrainingMode.EVAL,  # 推理模式
)

# 方案 1：重写为控制流算子
# 改 if 为 torch.cond，for 为 loop.unroll
from torch.onnx import export

# 方案 2：先用 torch.jit.script 转 TorchScript 再 export
scripted = torch.jit.script(model)
torch.onnx.export(scripted, dummy_input, "model.onnx", ...)

# 方案 3：自定义 onnx-symbolic
class MyOp(torch.autograd.Function):
    @staticmethod
    def symbolic(g, x):
        return g.op("custom::MyOp", x)  # 映射到自定义 ONNX 算子

import onnx
import onnxruntime as ort
import numpy as np

# 1. 检查模型结构
model = onnx.load("model.onnx")
onnx.checker.check_model(model)  # 拓扑正确性

# 2. 数值一致性验证
ort_session = ort.InferenceSession("model.onnx", providers=['CUDAExecutionProvider'])
ort_inputs = {k: v.cpu().numpy() for k, v in inputs.items()}
ort_output = ort_session.run(None, ort_inputs)[0]

torch_output = model(**inputs).logits.cpu().numpy()
np.testing.assert_allclose(ort_output, torch_output, rtol=1e-3, atol=1e-4)

# === 模型定义 ===
# PyTorch
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(784, 10)
    def forward(self, x):
        return self.fc(x)

# TensorFlow
class Net(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.fc = layers.Dense(10)
    def call(self, x):
        return self.fc(x)

# === 训练循环 ===
# PyTorch（显式）
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()

# TF（Keras fit 或 GradientTape）
@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        loss = criterion(model(x), y)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

# 1. torch.compile — 一行代码加速 1.5-3x
model = torch.compile(model, mode='reduce-overhead', fullgraph=True)

# 2. torch.export — AOT 编译，替代 torch.jit.trace
exported = torch.export.export(model, (dummy_input,))

# 3. FSDP v2 — 改进通信与显存
# 4. FlexAttention — 自定义注意力无 Triton 经验
# 5. torchao — 原生量化与稀疏化

# 基础用法
model = torch.compile(model)

# 指定模式
model = torch.compile(model, mode='default')      # 平衡编译时间和运行效率
model = torch.compile(model, mode='reduce-overhead')  # 减少 Python 开销，CUDA Graph
model = torch.compile(model, mode='max-autotune')  # 最大化性能，时间最长
model = torch.compile(model, mode='max-autotune-no-cudagraphs')  # 不使用 CUDA Graph

# fullgraph=True：禁止 Graph Break
model = torch.compile(model, fullgraph=True)

# 单独编译某个函数
@torch.compile
def my_function(x, y):
    return torch.relu(x @ y + 1)

# 1. 查看 Graph Break 位置
import torch._dynamo
torch._dynamo.config.verbose = True
torch._dynamo.config.suppress_errors = False  # 失败时抛错而非回退

# 2. 单独回退某个函数
@torch._dynamo.disable
def untraceable_function(x):
    ...

# 3. 查看编译后的 FX Graph
from torch._dynamo.utils import counters
print(counters)  # 统计 graph break 次数

# 4. 导出 FX Graph 调试
fx_graph = torch._dynamo.export(model)(dummy_input)
fx_graph.print_readable()

# 简单 MLP 训练
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(1024, 1024), nn.ReLU(),
            nn.Linear(1024, 1024), nn.ReLU(),
            nn.Linear(1024, 10)
        )
    def forward(self, x):
        return self.layers(x)

# RTX 4090, batch=4096, FP16
# 训练吞吐量
# Eager:           ~180k samples/s
# torch.compile:   ~340k samples/s  # 1.9x
# compile + AMP:   ~420k samples/s  # 2.3x

# 创建 Stream
stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()

# Stream 间同步
event = torch.cuda.Event()
with torch.cuda.stream(stream1):
    tensor1 = compute_a()  # 在 stream1 上
event.record(stream1)        # 记录事件
stream2.wait_event(event)    # stream2 等待
with torch.cuda.stream(stream2):
    tensor2 = compute_b(tensor1)  # 自动同步

# 设备级同步（阻塞当前设备）
torch.cuda.synchronize()

# 模式 1：数据预取与计算重叠
class PrefetchLoader:
    def __init__(self, loader, device):
        self.loader = loader
        self.device = device
        self.stream = torch.cuda.Stream()
        self.next_batch = None

    def __iter__(self):
        loader_iter = iter(self.loader)
        self._preload(loader_iter)
        return self

    def _preload(self, loader_iter):
        try:
            self.next_batch = next(loader_iter)
        except StopIteration:
            self.next_batch = None
            return
        with torch.cuda.stream(self.stream):
            self.next_batch = {
                k: v.to(self.device, non_blocking=True)
                for k, v in self.next_batch.items()
            }
        self.stream.synchronize()  # 必须 sync，否则竞争

    def __next__(self):
        batch = self.next_batch
        if batch is None:
            raise StopIteration
        self._preload(loader_iter)  # 后台预取下一个 batch
        return batch

# DDP 中梯度 bucket 内 AllReduce 与反向计算重叠
# 由 DDP 自动管理，无需手动
# FSDP 中 AllGather 与前向/反向计算重叠
# 由 FSDP 调度，无需手动

# 适合小模型、固定 shape、高频调用
g = torch.cuda.CUDAGraph()
# 预热
for _ in range(3):
    static_output = model(static_input)
# 捕获
with torch.cuda.graph(g):
    static_output = model(static_input)
# 复用
for input_batch in loader:
    static_input.copy_(input_batch)
    g.replay()
    # 复用 static_output

from torch.profiler import profile, ProfilerActivity, record_function, schedule

# 1. schedule：控制 profiler 范围
sched = schedule(wait=2, warmup=3, active=5, repeat=1)

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=sched,
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/prof'),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, (x, y) in enumerate(loader):
        with record_function("forward"):
            y_pred = model(x)
        with record_function("loss_backward"):
            loss = criterion(y_pred, y)
            loss.backward()
        optimizer.step()
        prof.step()

# 2. 打印关键统计
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

# 3. TensorBoard 可视化
# tensorboard --logdir=./log/prof

# 安装
pip install torchserve torch-model-archiver

# 打包模型
torch-model-archiver \
  --model-name resnet50 \
  --version 1.0 \
  --model-file model.py \
  --serialized-file resnet50.pt \
  --handler image_classifier

mkdir model_store && mv resnet50.mar model_store/

# 启动服务
torchserve --start \
  --model-store model_store \
  --models resnet50=resnet50.mar \
  --ncs        # 启用快照集成

# 1. 导出 ONNX / TorchScript / TF SavedModel
torch.onnx.export(model, dummy, "model.onnx",
                  dynamic_axes={'input': {0: 'batch'}})

# 2. 配置 model_repository/
# model_repository/my_model/config.pbtxt
"""
name: "my_model"
platform: "onnxruntime_onnx"
max_batch_size: 64
input [
  {
    name: "input"
    data_type: TYPE_FP32
    dims: [3, 224, 224]
  }
]
output [
  {
    name: "output"
    data_type: TYPE_FP32
    dims: [1000]
  }
]
instance_group [{ kind: KIND_GPU, count: 2 }]
dynamic_batching { preferred_batch_size: [16, 32] }
"""

# 3. 启动
docker run --gpus all -p 8000:8000 -p 8001:8001 -p 8002:8002 \
  -v $PWD/model_repository:/models nvcr.io/nvidia/tritonserver:24.08-py3 \
  tritonserver --model-repository=/models

import tritonclient.http as httpclient

triton_client = httpclient.InferenceServerClient(url="localhost:8000")
inputs = [httpclient.InferInput("input", [1, 3, 224, 224], "FP32")]
inputs[0].set_data_from_numpy(input_array)
outputs = [httpclient.InferRequestedOutput("output")]
response = triton_client.infer("my_model", inputs, outputs=outputs)
result = response.as_numpy("output")