kallewoof/tftf 来源：https://github.com/kallewoof/tftf

tftf (Transforming Transformers) —— 对 HuggingFace .safetensors 模型执行流式操作。张量逐个处理。整个模型永远不会加载到 RAM 或 VRAM —— 每个张量依次经过加载、变换、写入和释放，然后才处理下一个张量。

特性

安装

pip install -e .

要求 Python ≥ 3.11, PyTorch ≥ 2.0。FP8 反量化需要 PyTorch ≥ 2.1（用于 torch.float8_e4m3fn）。

CLI 参考

所有命令都接受单个 .safetensors 文件、包含 model.safetensors.index.json 的目录，或直接索引文件。所有写入命令都接受 --dry-run（仅验证不写入）、--sharded（写入分片文件 + index.json）和 --max-shard-size（每个分片的字节数，默认 5 GiB）。

info —— 查看模型

tftf info ./llama-7b/
tftf info ./model.safetensors --filter q_proj
tftf info ./DeepSeek-V3/ --dtype-summary

passthrough —— 不加载完整模型进行复制

# 复制并转换为 bfloat16
tftf passthrough -i ./model.safetensors -o ./model-bf16.safetensors --dtype bfloat16

# 复制分片模型，输出为新分片
tftf passthrough -i ./llama-70b/ -o ./llama-70b-copy/ --sharded

# 仅复制注意力权重
tftf passthrough -i ./model.safetensors -o ./attn.safetensors \
    --include '*self_attn*'

# 仅验证不写入
tftf passthrough -i ./model.safetensors -o /dev/null --dry-run

dequant-fp8 —— 将细粒度 FP8 反量化为 BF16/FP16

# 将 DeepSeek-V3 反量化为 bfloat16，输出分片
tftf dequant-fp8 \
    -i ./DeepSeek-V3/ \
    -o ./DeepSeek-V3-bf16/ \
    --sharded

# 一键反量化并融合 LoRA 适配器
tftf dequant-fp8 \
    -i ./DeepSeek-V3/ \
    -o ./merged/ \
    --dtype bfloat16 \
    --merge-lora ./my-lora/adapter_model.safetensors \
    --sharded

# 试运行，在提交磁盘空间之前验证流水线
tftf dequant-fp8 \
    -i ./DeepSeek-V3/ \
    -o /dev/null \
    --dtype bfloat16 \
    --dry-run

merge-lora —— 将 LoRA 适配器融合到基座模型

tftf merge-lora \
    -b ./llama-7b/ \
    -a ./my-lora/adapter_model.safetensors \
    -o ./merged.safetensors \
    --dtype bfloat16

# 输出为分片并重命名键
tftf merge-lora \
    -b ./llama-7b/ \
    -a ./adapter_model.safetensors \
    -o ./merged/ --sharded \
    --rename '^transformer\.h\.' 'model.layers.'

merge-fsdp-lora —— 融合 per-rank FSDP 分片的 LoRA

# 显式指定 per-rank 文件
tftf merge-fsdp-lora \
    -b ./llama-7b/ \
    -s ./run/rank_00.safetensors \
    -s ./run/rank_01.safetensors \
    -o ./merged.safetensors

# 分片文件目录（按字母排序 = rank 顺序）
tftf merge-fsdp-lora \
    -b ./llama-7b/ --shard-dir ./run/ \
    -o ./merged/ --sharded --dtype bfloat16

validate —— 不写入的验证

tftf validate ./model.safetensors
tftf validate ./llama-70b/
tftf validate ./model.safetensors --pipe dtype:bfloat16
tftf validate ./model.safetensors --pipe filter:*q_proj*

Python API

一键完成 FP8 反量化 + LoRA 融合

from tftf import (
    Pipeline,
    ShardedSafetensorsReader,
    ShardedWriter,
    FP8DequantPipe,
    LoRAMergePipe,
)
import torch

pipe = FP8DequantPipe(torch.bfloat16) | LoRAMergePipe("./my-lora/adapter_model.safetensors")

Pipeline(
    reader=ShardedSafetensorsReader.from_path("./DeepSeek-V3/"),
    pipe=pipe,
    writer=ShardedWriter("./DeepSeek-V3-bf16-merged/"),
).run()

试运行验证

from tftf import Pipeline, SafetensorsReader, NullWriter, FP8DequantPipe
import torch

writer = NullWriter()
Pipeline(
    reader=SafetensorsReader("./model.safetensors"),
    pipe=FP8DequantPipe(torch.bfloat16),
    writer=writer,
).run()
print(writer.report.summary())
assert writer.report.ok

LoRA 融合并转换数据类型

from tftf import Pipeline, SafetensorsReader, StreamingWriter, LoRAMergePipe, DTypeCastPipe
import torch

pipe = LoRAMergePipe("adapter_model.safetensors") | DTypeCastPipe(torch.float16)

Pipeline(
    reader=SafetensorsReader("model.safetensors"),
    pipe=pipe,
    writer=StreamingWriter("merged.safetensors"),
).run()

架构

model_pipe/
├── cli.py                  Click CLI (info, passthrough, dequant-fp8,
│                           merge-lora, merge-fsdp-lora, validate)
├── pipeline.py             两遍协调器
├── pipes/
│   ├── base.py             Pipe (ABC), CompoundPipe, TensorRecord, TensorMeta
│   ├── passthrough.py      恒等管道
│   ├── dtype_cast.py       DTypeCastPipe
│   ├── key_filter.py       KeyFilterPipe (glob 包含/排除)
│   ├── key_rename.py       KeyRenamePipe (正则替换)
│   ├── _lora_base.py       LoRAMergeBase (共享融合逻辑)
│   ├── lora_merge.py       LoRAMergePipe (单个适配器文件)
│   ├── fsdp_lora_merge.py  FSDPShardMergePipe (per-rank 分片)
│   └── fp8_dequant.py      FP8DequantPipe (细粒度 FP8 → BF16/FP16)
├── io/
│   ├── reader.py           SafetensorsReader (单文件, mmap)
│   ├── sharded_reader.py   ShardedSafetensorsReader (index.json)
│   ├── writer.py           StreamingWriter (单文件输出)
│   ├── sharded_writer.py   ShardedWriter (多分片输出)
│   └── null_writer.py      NullWriter + ValidationReport (试运行)
└── utils/
    ├── lora.py             LoRA 键映射、融合数学
    ├── fsdp.py             FSDP 分片发现与重构
    └── fp8.py              FP8 数据类型辅助、向量化反量化

两遍流水线

safetensors 格式要求完整的头部（所有张量的名称、形状、数据类型和字节偏移量）位于文件开头，在数据之前。这给流式输出带来了“先有鸡还是先有蛋”的问题，通过两遍解决：

阶段 1 — 元数据扫描（不加载张量数据）
──────────────────────────────────────────────────
reader.iter_meta() 仅读取 JSON 头部（mmap 头部区域）
      ↓
pipe.process_meta() 转换键/形状/数据类型声明
      ↓
writer.prepare(metas) 将 safetensors 头部写入磁盘

阶段 2 — 数据流（RAM 中一次只有一个张量）
──────────────────────────────────────────
reader.iter_records() mmap：逐页调入一个张量
      ↓
pipe.process() 变换：反量化 / 融合 / 转换 / 过滤
      ↓
writer.write_record() 追加写入原始张量字节
      ↓
del record          GC 在加载下一个张量前回收
      (重复)
      ↓
writer.finalize()

Pipe 接口

class Pipe(ABC):
    def process(self, records: Iterator[TensorRecord]) -> Iterator[TensorRecord]:
        ...

    def process_meta(self, metas: Iterator[TensorMeta]) -> Iterator[TensorMeta]:
        ...

    def setup(self) -> None:
        ...  # 在 process() 之前调用一次

    def teardown(self) -> None:
        ...  # 在 process() 之后调用一次

    def __or__(self, other) -> CompoundPipe:
        ...

    def __repr__(self) -> str:
        ...

实现 process() 的规则：

• 要懒惰 —— 消费一条记录，yield 零或多条，立即释放
• 不要缓冲整个流
• 可以丢弃记录（过滤）或添加记录（注入）
• 尽可能快地释放张量：yield 之后执行 del record.tensor

仅在管道改变键、形状或数据类型时覆盖 process_meta()。默认是恒等操作。

编写新管道

from tftf.pipes.base import Pipe, TensorRecord, TensorMeta
from typing import Iterator
import torch

class MyQuantisePipe(Pipe):
    """示例：将每个 float32 权重量化为 int8。"""

    def process_meta(self, metas: Iterator[TensorMeta]) -> Iterator[TensorMeta]:
        for meta in metas:
            new_dtype = torch.int8 if meta.dtype == torch.float32 else meta.dtype
            yield TensorMeta(key=meta.key, dtype=new_dtype, shape=meta.shape)

    def process(self, records: Iterator[TensorRecord]) -> Iterator[TensorRecord]:
        for record in records:
            t = record.tensor.to(torch.int8) if record.tensor.dtype == torch.float32 \
                else record.tensor
            yield TensorRecord(key=record.key, tensor=t)
            del t

    def __repr__(self) -> str:
        return "MyQuantisePipe()"

FP8 格式细节

像 DeepSeek-V3 这样的模型将量化后的权重存储为：

反量化逐块重构全精度权重：

W_out[r:r+128, c:c+128] = W_fp8[r:r+128, c:c+128].float() * scale_inv[r//128, c//128]

FP8DequantPipe 使用向量化的广播乘法（填充 → reshape → permute → 乘 → 去填充），而不是 Python 循环，对于大矩阵速度提高了几个数量级。缩放张量会被自动检测并从输出中丢弃。

LoRA 键映射

LoRAMergePipe 和 FSDPShardMergePipe 自动检测 PEFT 命名约定：

adapter_config.json 会从适配器目录中自动检测，用于获取 r 和 lora_alpha。

运行测试

pip install -e ".[dev]"
pytest -v

4 个测试文件共 91 个测试。所有测试均在不下载真实模型的情况下运行 —— 在进程中生成合成张量。

要求

• Python ≥ 3.11
• PyTorch ≥ 2.0（≥ 2.1 以支持 FP8）
• safetensors ≥ 0.4
• click ≥ 8.1
• tqdm ≥ 4.60