用 GPUDirect + FSx for Lustre 加速百亿参数 LLM 加载，TurboQuant 扩展上下文窗口

[持久存储] → CPU RAM (DMA) → GPU HBM (PCIe copy)

[FSx for Lustre] → GPU HBM (RDMA, bypass CPU)

# 创建 S3 桶存放模型权重（如果还没有）
aws s3api create-bucket \
  --bucket llm-models-repo \
  --region us-east-1

# 上传模型到 S3（假设已量化为 TurboQuant 格式）
aws s3 cp ./models/llama-70b-tq4/ s3://llm-models-repo/llama-70b-tq4/ --recursive

# 创建 FSx for Lustre，关联 S3 桶作为持久层
aws fsx create-file-system \
  --file-system-type LUSTRE \
  --subnet-id subnet-0abc123def456 \
  --security-group-ids sg-0xyz789 \
  --lustre-configuration DeploymentType=PERSISTENT_2,PerUnitStorageThroughput=125,AutoImportPolicy=NEW_CHANGED_DELETED,ImportPath=s3://llm-models-repo/ \
  --storage-capacity 2400 \
  --region us-east-1

# 安装 Lustre 客户端
sudo yum install -y lustre-client

# 获取 FSx 挂载点（从创建返回的 DNS 名称）
MOUNT_DNS="fs-0123456.fsx.us-east-1.amazonaws.com"

# 挂载文件系统
sudo mkdir -p /fsx/llm-models
sudo mount -t lustre ${MOUNT_DNS}@tcp:/fsx /fsx/llm-models

# 验证模型文件可见
ls -lh /fsx/llm-models/llama-70b-tq4/
# 应看到量化后的权重文件，总大小约 35–40 GB（4-bit 量化后）

# 安装 NVIDIA MLNX_OFED 并启用 GPUDirect Storage
sudo apt-get install -y mlnx-ofed-driver
sudo modprobe nvidia-peermem
# 验证 GDS 可用
nvidia-smi -q | grep "GPUDirect"

"""
使用 GDS 直读路径加载 TurboQuant 量化模型到 GPU
依赖：transformers, turboquant, nvidia-gds-binding
"""
import os
import time
from pathlib import Path

# 关键：让 PyTorch 使用 GDS 直读路径
# 环境变量启用 CUDA 的 GDS 功能
os.environ["CUDA_GDS"] = "1"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_PATH = "/fsx/llm-models/llama-70b-tq4"

def load_model_with_gds():
    """通过 GDS 直读路径加载量化模型"""
    start = time.time()

    # TurboQuant 格式模型会自动在 GPU 上解量化
    # GDS 确保数据从 Lustre 直达 GPU HBM
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.float8_e4m3fn,  # H100 推理推荐 FP8
        device_map="auto",                 # 自动分配到 8 张 H100
        trust_remote_code=True,
    )

    load_time = time.time() - start
    # 打印每张卡的显存占用
    for i in range(torch.cuda.device_count()):
        mem = torch.cuda.memory_allocated(i) / 1024**3
        print(f"GPU {i}: {mem:.1f} GB HBM used")

    print(f"模型加载耗时: {load_time:.1f}s")
    return model, tokenizer

model, tokenizer = load_model_with_gds()

# 检查 KV Cache 可用空间——量化后腾出的 HBM 让上下文窗口更大
# 70B 模型 FP16 需 ~140 GB HBM，4-bit 量化 + FP8 推理约 35-40 GB
# 8×H100 总 HBM 640 GB，剩余约 600 GB 可用于 KV Cache
# 理论上支持 128K+ 上下文窗口

# 传统路径（关闭 GDS，数据经 CPU 中转）
CUDA_GDS=0 python load_benchmark.py --model /fsx/llm-models/llama-70b-tq4

# GDS 直读路径
CUDA_GDS=1 python load_benchmark.py --model /fsx/llm-models/llama-70b-tq4

# 预期结果（p5.48xlarge + Lustre 500 throughput）：
# 传统路径: ~25-30s（70B 量化模型）
# GDS 路径:  ~10-15s
# 提速约 40-50%