kallewoof/tftf

Source: https://github.com/kallewoof/tftf

tftf (Transforming Transformers)

Streaming operations on HuggingFace .safetensors models.

Tensors are processed one at a time. The full model is never loaded into RAM or VRAM — each tensor is loaded, transformed, written, and freed before the next one is touched.

Features

Install

pip install -e .

Requires Python ≥ 3.11, PyTorch ≥ 2.0. FP8 dequantisation requires PyTorch ≥ 2.1 (for torch.float8_e4m3fn).

CLI reference

All commands accept single .safetensors files, directories containing model.safetensors.index.json, or index files directly.

All write commands accept --dry-run (validate without writing), --sharded (write as shard files + index.json), and --max-shard-size (bytes per shard, default 5 GiB).

info — inspect a model

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

passthrough — copy without loading the full model

# Copy and cast to bfloat16
tftf passthrough -i ./model.safetensors -o ./model-bf16.safetensors --dtype bfloat16

# Copy a sharded model, writing output as new shards
tftf passthrough -i ./llama-70b/ -o ./llama-70b-copy/ --sharded

# Copy only attention weights
tftf passthrough -i ./model.safetensors -o ./attn.safetensors \
    --include '*self_attn*'

# Validate without writing
tftf passthrough -i ./model.safetensors -o /dev/null --dry-run

dequant-fp8 — dequantise fine-grained FP8 to BF16/FP16

# Dequantise DeepSeek-V3 to bfloat16, writing sharded output
tftf dequant-fp8 \
    -i ./DeepSeek-V3/ \
    -o ./DeepSeek-V3-bf16/ \
    --sharded

# Dequantise then fuse a LoRA adapter in one pass
tftf dequant-fp8 \
    -i ./DeepSeek-V3/ \
    -o ./merged/ \
    --dtype bfloat16 \
    --merge-lora ./my-lora/adapter_model.safetensors \
    --sharded

# Dry-run to validate the pipeline before committing disk space
tftf dequant-fp8 \
    -i ./DeepSeek-V3/ \
    -o /dev/null \
    --dtype bfloat16 \
    --dry-run

merge-lora — fuse a LoRA adapter into a base model

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

# Write as shards with key renaming
tftf merge-lora \
    -b ./llama-7b/ \
    -a ./adapter_model.safetensors \
    -o ./merged/ --sharded \
    --rename '^transformer\.h\.' 'model.layers.'

merge-fsdp-lora — fuse a per-rank FSDP-sharded LoRA

# Explicit per-rank files
tftf merge-fsdp-lora \
    -b ./llama-7b/ \
    -s ./run/rank_00.safetensors \
    -s ./run/rank_01.safetensors \
    -o ./merged.safetensors

# Directory of shard files (sorted alphabetically = rank order)
tftf merge-fsdp-lora \
    -b ./llama-7b/ --shard-dir ./run/ \
    -o ./merged/ --sharded --dtype bfloat16

validate — validate without writing

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 dequantisation + LoRA merge in one pass

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()

Dry-run validation

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 merge with dtype cast

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()

Architecture

model_pipe/
├── cli.py                  Click CLI (info, passthrough, dequant-fp8,
│                           merge-lora, merge-fsdp-lora, validate)
├── pipeline.py             Two-pass orchestrator
├── pipes/
│   ├── base.py             Pipe (ABC), CompoundPipe, TensorRecord, TensorMeta
│   ├── passthrough.py      Identity pipe
│   ├── dtype_cast.py       DTypeCastPipe
│   ├── key_filter.py       KeyFilterPipe (glob include/exclude)
│   ├── key_rename.py       KeyRenamePipe (regex substitution)
│   ├── _lora_base.py       LoRAMergeBase (shared merge logic)
│   ├── lora_merge.py       LoRAMergePipe (single adapter file)
│   ├── fsdp_lora_merge.py  FSDPShardMergePipe (per-rank shards)
│   └── fp8_dequant.py      FP8DequantPipe (fine-grained FP8 → BF16/FP16)
├── io/
│   ├── reader.py           SafetensorsReader (single file, mmap)
│   ├── sharded_reader.py   ShardedSafetensorsReader (index.json)
│   ├── writer.py           StreamingWriter (single file output)
│   ├── sharded_writer.py   ShardedWriter (multi-shard output)
│   └── null_writer.py      NullWriter + ValidationReport (dry-run)
└── utils/
    ├── lora.py             LoRA key mapping, merge math
    ├── fsdp.py             FSDP shard discovery and reconstruction
    └── fp8.py              FP8 dtype helpers, vectorised dequantisation

Two-pass pipeline

The safetensors format requires the complete header (all tensor names, shapes, dtypes, and byte offsets) at the start of the file, before any data. This creates a chicken-and-egg problem for streaming output, solved with two passes:

Phase 1 — metadata scan  (no tensor data loaded)
──────────────────────────────────────────────────
reader.iter_meta()         reads JSON header only (mmap header section)
    ↓
pipe.process_meta()        transforms key/shape/dtype declarations
    ↓
writer.prepare(metas)      writes safetensors header to disk

Phase 2 — data stream  (one tensor in RAM at a time)
──────────────────────────────────────────────────────
reader.iter_records()      mmap: pages in one tensor at a time
    ↓
pipe.process()             transform: dequant / merge / cast / filter
    ↓
writer.write_record()      appends raw tensor bytes
    ↓  del record          GC reclaims before loading next tensor
    (repeat)
    ↓
writer.finalize()

Pipe interface

class Pipe(ABC):
    def process(self, records: Iterator[TensorRecord]) -> Iterator[TensorRecord]: ...
    def process_meta(self, metas: Iterator[TensorMeta]) -> Iterator[TensorMeta]: ...
    def setup(self) -> None: ...     # called once before process()
    def teardown(self) -> None: ...  # called once after process()
    def __or__(self, other) -> CompoundPipe: ...
    def __repr__(self) -> str: ...

Rules for implementing process():

• Be lazy — consume one record, yield zero-or-more, free immediately
• Don’t buffer the whole stream
• It is valid to drop records (filter) or add records (inject)
• Free tensors as soon as possible: del record.tensor after yielding

Override process_meta() only when your pipe changes keys, shapes, or dtypes. The default is the identity.

Writing a new pipe

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

class MyQuantisePipe(Pipe):
    """Example: quantise every float32 weight to 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 format details

Models like DeepSeek-V3 store quantised weights as:

Dequantisation reconstructs the full-precision weight block by block:

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

FP8DequantPipe uses a vectorised broadcast-multiply (pad → reshape → permute → multiply → unpad) rather than a Python loop, making it orders of magnitude faster for large matrices. Scale tensors are automatically detected and dropped from output.

LoRA key mapping

LoRAMergePipe and FSDPShardMergePipe auto-detect PEFT naming conventions:

adapter_config.json is auto-detected from the adapter directory for r and lora_alpha.

Running tests

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

91 tests across 4 test files. All tests run without downloading any real model — synthetic tensors are generated in-process.

Requirements

• Python ≥ 3.11
• PyTorch ≥ 2.0 (≥ 2.1 for FP8 support)
• safetensors ≥ 0.4
• click ≥ 8.1
• tqdm ≥ 4.60