Skip to content

Commit 31948fb

Browse files
HosseinKaviani-HHossein Kavianihamedani
HosseinKaviani-H
and
Hossein Kavianihamedani
authored
Enable SFT for HF Transformer backed trainer (#3243)
## Summary Enables supervised fine-tuning (SFT) for any HuggingFace model in the transformers_modeling_backend experiment, building on the native SFT infrastructure from #2556. - Uses FlexAttention with BlockMask for packed sequences. Builds the BlockMask in `post_dataloading_process` (before parallelism sharding) using TorchTitan's existing `get_causal_mask_mod` + `get_document_mask_mod` — same mask infrastructure as native TorchTitan SFT. - `HFFlexAttention` subclasses TorchTitan's `FlexAttention` kernel module, handling HF's tensor layout via transposes. Registered as a submodule on each HF attention layer so `apply_cp_to_forward` works naturally — no globals needed. - Registers `flex_torchtitan` as a custom attention implementation that bypasses HF's internal mask creation (which builds at the wrong point for TP+SP) and routes through a compile-friendly `flex_attention` call. - `SFTTrainer` subclass overrides `post_dataloading_process` to build the BlockMask at full sequence length before any parallelism touches the data. - `HFTransformerStateDictAdapter` enables loading pretrained HF weights (trivial `model.` prefix mapping, handles weight tying). - `HFBackendTokenizer` passes `bos_token`/`eos_token` to chat templates and fixes an edge case where BOS and EOS share the same token string. - TitanDenseModelConfig defaults changed to None so HF config values aren't overridden when using full-size models. - No dynamo `cache_size_limit` workaround needed: relies on the upstream nested-`AsyncCollectiveTensor` fix (pytorch#186442), so the flex path stays under the default recompile limit across all parallelisms. - All changes within `torchtitan/experiments/transformers_modeling_backend/` — no core modifications. ## Files changed | File | Description | |---|---| | trainer.py (new) | `SFTTrainer` builds BlockMask before parallelism sharding, handles CP load balancer | | state_dict_adapter.py (new) | `model.` prefix strip/add for HF weight loading, handles weight tying | | tokenizer.py (new) | Passes special tokens to chat templates, fixes shared BOS/EOS edge case | | model.py | `HFFlexAttention` subclass + `_flex_torchtitan_attention_forward`; `forward()` accepts positions/attention_masks; guard `intermediate_size`/`head_dim` recalculation; disable `attention_dropout` for FlexAttention | | configs.py | `build()` override to use `SFTTrainer` | | parallelize.py | Collects `HFFlexAttention` modules and calls `apply_cp_to_forward` for CP | | __init__.py | `sft_debugmodel`/`sft_full` flavors with `attn_implementation="flex_torchtitan"`; wire `state_dict_adapter`; None defaults for TitanDenseModelConfig | | config_registry.py | SFT config functions (random init + pretrained weight loading) using `ChatDataLoader` and `HFBackendTokenizer` | | tests/integration_tests.py | 2-GPU SFT integration test | ## Parallelism support | Parallelism | Status | |---|---| | FSDP | Works | | FSDP + TP | Works | | FSDP + TP + PP | Works | | FSDP + AC (selective/full) | Works | | FSDP + compile | Works | | FSDP + TP + compile | Works | | FSDP + CP | Works | | CP + compile | Works | | TP + CP | Works | | FSDP + TP + CP | Works | | FSDP + TP + CP + compile | Works | ## Test plan - [x] Pretraining regression: `transformers_modeling_backend_debugmodel` identical loss before/after - [x] SFT debugmodel (Qwen3 arch): loss 11.0 → 3.3, 2 GPUs - [x] SFT debugmodel (Mistral arch): loss 10.9 → 8.3 — architecture-agnostic - [x] SFT with pretrained Qwen3-0.6B: loss 13.4 → 6.1 - [x] SFT with pretrained Qwen2.5-0.5B-Instruct: loss 5.3 → 1.7 - [x] SFT with pretrained Llama-3.2-1B-Instruct: loss 7.6 → 0.0001 - [x] SFT with pretrained Seed-Coder-8B-Instruct: loss 6.8 → 2.1 - [x] SFT with pretrained Granite-3.1-2B-Instruct: loss 50.2 → 12.2 - [x] FSDP + TP (4 GPU): works - [x] FSDP + TP + PP (8 GPU): works - [x] FSDP + compile (2 GPU): works - [x] FSDP + TP + compile (4 GPU): works - [x] FSDP + CP (2 GPU): works - [x] CP + compile (2 GPU): works - [x] TP + CP (4 GPU): works - [x] FSDP + TP + CP (8 GPU): works - [x] FSDP + TP + CP + compile (8 GPU, full 28-layer qwen3_0.6b): works — 0 cache-limit hits at default recompile limit - [x] Reproducibility: compare loss curves against native SFT on same model/data - [x] SFT loss/KL parity vs native (#2556), qwen3_0.6b, same weights + packed batch: max |loss diff| 3.4e-06, max KL 5.0e-09 Numeric verification results: ### Native TorchTitan vs HuggingFace — logit parity (qwen3_0.6b, same weights + input, fp32, 161 positions) | metric | value | |---|---| | top1 match | 100.0000% | | top5 overlap | 100.0000% | | KL(HF ‖ native) | 1.60e-06 | | KL(native ‖ HF) | 6.78e-07 | | max \|logit diff\| | 9.06e-05 | | mean \|logit diff\| | 5.14e-06 | | cosine similarity | 1.000000 | --------- Co-authored-by: Hossein Kavianihamedani <hosseinkh@fb.com>
1 parent 1d2d58c commit 31948fb

8 files changed

Lines changed: 393 additions & 17 deletions

File tree

‎torchtitan/experiments/transformers_modeling_backend/__init__.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,9 +59,25 @@ class TitanDenseModelConfig:
5959
n_kv_heads=16,
6060
),
6161
),
62+
"sft_debugmodel": HFTransformerModel.Config(
63+
titan_dense_config=TitanDenseModelConfig(
64+
dim=256,
65+
n_layers=2,
66+
n_heads=16,
67+
n_kv_heads=16,
68+
attn_mask_type="block_causal",
69+
),
70+
attn_implementation="flex_torchtitan",
71+
),
6272
"full": HFTransformerModel.Config(
6373
titan_dense_config=TitanDenseModelConfig(),
6474
),
75+
"sft_full": HFTransformerModel.Config(
76+
titan_dense_config=TitanDenseModelConfig(
77+
attn_mask_type="block_causal",
78+
),
79+
attn_implementation="flex_torchtitan",
80+
),
6581
}
6682

6783

‎torchtitan/experiments/transformers_modeling_backend/config_registry.py‎

Lines changed: 100 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,14 @@
1414
from torchtitan.experiments.transformers_modeling_backend.configs import (
1515
TransformersBackendConfig,
1616
)
17-
from torchtitan.hf_datasets.text_datasets import HuggingFaceTextDataLoader
17+
from torchtitan.hf_datasets.text_datasets import (
18+
ChatDataLoader,
19+
HuggingFaceTextDataLoader,
20+
)
1821
from torchtitan.tools.profiler import Profiler
1922

2023
from . import model_registry
24+
from .tokenizer import HFBackendTokenizer
2125

2226

2327
def transformers_modeling_backend_debugmodel() -> TransformersBackendConfig:
@@ -81,3 +85,98 @@ def transformers_modeling_backend_full() -> TransformersBackendConfig:
8185
),
8286
activation_checkpoint=SelectiveAC.Config(),
8387
)
88+
89+
90+
def transformers_modeling_backend_sft_full() -> TransformersBackendConfig:
91+
"""SFT config with real HF pretrained weights loaded via initial_load_in_hf."""
92+
93+
def process_sample(sample):
94+
return [
95+
{"role": "user", "content": sample["question"]},
96+
{"role": "assistant", "content": sample["answer"]},
97+
]
98+
99+
return TransformersBackendConfig(
100+
loss=CrossEntropyLoss.Config(),
101+
hf_assets_path="./tests/assets/qwen3_0.6b",
102+
hf_model="Qwen/Qwen3-0.6B",
103+
model_spec=model_registry("sft_full"),
104+
tokenizer=HFBackendTokenizer.Config(),
105+
optimizer=default_adamw(lr=2e-5),
106+
lr_scheduler=LRSchedulersContainer.Config(
107+
warmup_steps=2,
108+
decay_ratio=0.8,
109+
decay_type="linear",
110+
min_lr_factor=0.0,
111+
),
112+
training=TrainingConfig(
113+
local_batch_size=2,
114+
seq_len=2048,
115+
steps=10,
116+
),
117+
dataloader=ChatDataLoader.Config(
118+
dataset_path="json",
119+
load_dataset_kwargs={
120+
"data_files": "tests/assets/sft_test/data.json",
121+
"split": "train",
122+
},
123+
sample_processor=process_sample,
124+
),
125+
metrics=MetricsProcessor.Config(log_freq=1),
126+
checkpoint=CheckpointManager.Config(
127+
enable=True,
128+
initial_load_in_hf=True,
129+
initial_load_model_only=True,
130+
interval=10,
131+
last_save_model_only=False,
132+
),
133+
activation_checkpoint=SelectiveAC.Config(),
134+
)
135+
136+
137+
def transformers_modeling_backend_sft_debugmodel() -> TransformersBackendConfig:
138+
"""SFT debug config for the transformers backend using ChatDataLoader."""
139+
140+
def process_sample(sample):
141+
return [
142+
{"role": "user", "content": sample["question"]},
143+
{"role": "assistant", "content": sample["answer"]},
144+
]
145+
146+
return TransformersBackendConfig(
147+
loss=CrossEntropyLoss.Config(),
148+
hf_assets_path="./tests/assets/tokenizer",
149+
hf_model="Qwen/Qwen3-4B-Instruct-2507",
150+
model_spec=model_registry("sft_debugmodel"),
151+
tokenizer=HFBackendTokenizer.Config(),
152+
optimizer=default_adamw(lr=8e-4),
153+
lr_scheduler=LRSchedulersContainer.Config(
154+
warmup_steps=2,
155+
decay_ratio=0.8,
156+
decay_type="linear",
157+
min_lr_factor=0.0,
158+
),
159+
training=TrainingConfig(
160+
# Keep this small: this debug model uses the full Qwen3 vocab
161+
# (~152k), so cross-entropy materializes a
162+
# local_batch_size * seq_len * vocab logits tensor. batch=8,
163+
# seq=2048 is ~9GB in fp32 and OOMs the 22GB CI GPUs.
164+
local_batch_size=1,
165+
seq_len=1024,
166+
steps=10,
167+
),
168+
dataloader=ChatDataLoader.Config(
169+
dataset_path="json",
170+
load_dataset_kwargs={
171+
"data_files": "tests/assets/sft_test/data.json",
172+
"split": "train",
173+
},
174+
sample_processor=process_sample,
175+
),
176+
metrics=MetricsProcessor.Config(log_freq=1),
177+
checkpoint=CheckpointManager.Config(
178+
interval=10,
179+
last_save_model_only=False,
180+
),
181+
activation_checkpoint=SelectiveAC.Config(),
182+
)

‎torchtitan/experiments/transformers_modeling_backend/configs.py‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,10 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7-
from dataclasses import dataclass
7+
from torchtitan.experiments.transformers_modeling_backend.trainer import SFTTrainer
88

9-
from torchtitan.trainer import Trainer
10-
11-
12-
@dataclass(kw_only=True, slots=True)
13-
class TransformersBackendConfig(Trainer.Config):
14-
hf_model: str = ""
15-
"""HuggingFace model ID (e.g., 'Qwen/Qwen2.5-7B')"""
9+
# The HF-backend config is just SFTTrainer's own Config. SFTTrainer defines a
10+
# nested Config, so Configurable.__init_subclass__ auto-wires _owner=SFTTrainer;
11+
# the inherited Config.build() then constructs an SFTTrainer (and keeps the
12+
# config-field vs build()-kwarg overlap check) -- no reimplementation needed.
13+
TransformersBackendConfig = SFTTrainer.Config

‎torchtitan/experiments/transformers_modeling_backend/model.py‎

Lines changed: 123 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,99 @@
1212
import torch
1313
from torch import nn
1414
from torch.nn import init
15+
16+
from torch.nn.attention.flex_attention import BlockMask
17+
1518
from transformers import AutoConfig
1619
from transformers.configuration_utils import PretrainedConfig
1720
from transformers.integrations.sdpa_attention import sdpa_attention_forward
1821
from transformers.modeling_utils import AttentionInterface, PreTrainedModel
1922

23+
from torchtitan.models.common.attention import FlexAttention
2024
from torchtitan.models.utils import get_dense_model_nparams_and_flops
2125
from torchtitan.protocols.model import BaseModel
2226
from torchtitan.protocols.module import ModuleDict
2327
from torchtitan.tools.logging import logger
2428

2529

30+
# Shape suffix legend (matches torchtitan/models/common/attention.py):
31+
# B = batch, L = sequence length, N = num heads, H = head dimension.
32+
# HF attention layout is (B, N, L, H); native TorchTitan layout is (B, L, N, H).
33+
class HFFlexAttention(FlexAttention):
34+
"""FlexAttention kernel for HF models.
35+
36+
Accepts Q/K/V in native TorchTitan layout (B, L, N, H) so that
37+
apply_cp_to_forward's K/V all-gather (dim=1, the seq dim) works correctly.
38+
The caller (_flex_torchtitan_attention_forward) transposes from HF layout
39+
before calling this module, and transposes back after.
40+
41+
This subclass passes the isinstance(FlexAttention) check in
42+
apply_cp_to_forward, so CP wrapping works naturally.
43+
"""
44+
45+
def forward(
46+
self,
47+
q_BLNH: torch.Tensor,
48+
k_BLNH: torch.Tensor,
49+
v_BLNH: torch.Tensor,
50+
*,
51+
attention_masks: BlockMask | None = None,
52+
scale: float | None = None,
53+
enable_gqa: bool = True,
54+
**kwargs,
55+
) -> torch.Tensor:
56+
from torch.nn.attention.flex_attention import flex_attention
57+
58+
# Transpose to (B, N, L, H) for the flex_attention kernel
59+
q_BNLH = q_BLNH.transpose(1, 2)
60+
k_BNLH = k_BLNH.transpose(1, 2)
61+
v_BNLH = v_BLNH.transpose(1, 2)
62+
out_BNLH = flex_attention(
63+
q_BNLH,
64+
k_BNLH,
65+
v_BNLH,
66+
block_mask=attention_masks,
67+
scale=scale,
68+
enable_gqa=enable_gqa,
69+
)
70+
# Transpose back to (B, L, N, H)
71+
return out_BNLH.transpose(1, 2)
72+
73+
74+
def _flex_torchtitan_attention_forward(
75+
module, query_BNLH, key_BNLH, value_BNLH, attention_mask, **kwargs
76+
):
77+
"""FlexAttention forward registered via AttentionInterface.
78+
79+
Routes the kernel call through the HFFlexAttention module attached to
80+
each HF attention layer, so apply_cp_to_forward's K/V all-gather wrapping
81+
applies automatically.
82+
83+
Transposes Q/K/V from HF layout (B, N, L, H) to native TorchTitan layout
84+
(B, L, N, H) before calling the module, so CP's dim=1 all-gather targets
85+
the sequence dimension correctly.
86+
"""
87+
scaling = kwargs.get("scaling")
88+
block_mask = attention_mask if isinstance(attention_mask, BlockMask) else None
89+
90+
# This forward is only registered as the "flex_torchtitan" attention impl,
91+
# and HFTransformerModel attaches a _flex_kernel to every self_attn when that
92+
# impl is selected, so the kernel module is always present here.
93+
flex_module = module._flex_kernel
94+
95+
# HF layout (B, N, L, H) -> native layout (B, L, N, H)
96+
q_BLNH = query_BNLH.transpose(1, 2)
97+
k_BLNH = key_BNLH.transpose(1, 2)
98+
v_BLNH = value_BNLH.transpose(1, 2)
99+
out_BLNH = flex_module(
100+
q_BLNH, k_BLNH, v_BLNH, attention_masks=block_mask, scale=scaling
101+
)
102+
# HF attention interface expects output as (B, L, N, H) (batch, seq, heads,
103+
# head_dim) -- same layout sdpa_attention_forward returns. out_BLNH is
104+
# already in that layout, so return it directly (no transpose).
105+
return out_BLNH, None
106+
107+
26108
class SliceableModuleDict(ModuleDict):
27109
"""
28110
A ModuleDict that supports slicing like ModuleList.
@@ -178,10 +260,18 @@ def _configure_hf_attention(self, attn_implementation: str):
178260
"""Configure HuggingFace attention settings."""
179261
self._titan_injected_model_args["attn_implementation"] = attn_implementation
180262
self.attn_implementation = attn_implementation
181-
# NOTE:(3outeille):This will force create_causal_mask to return None
182-
AttentionInterface._global_mapping[
183-
attn_implementation
184-
] = sdpa_attention_forward
263+
if attn_implementation == "flex_torchtitan":
264+
AttentionInterface._global_mapping[
265+
attn_implementation
266+
] = _flex_torchtitan_attention_forward
267+
elif attn_implementation not in AttentionInterface._global_mapping:
268+
logger.info(
269+
f"'{attn_implementation}' not in HF AttentionInterface registry, "
270+
"defaulting to sdpa_attention_forward"
271+
)
272+
AttentionInterface._global_mapping[
273+
attn_implementation
274+
] = sdpa_attention_forward
185275

186276
def _create_getter_setter_dynamically(self, has_moe: bool):
187277
"""
@@ -264,6 +354,14 @@ def update_from_config(
264354
self.mlp_bias = False
265355
self.use_cache = False
266356
self.initializer_range = 1.0 # use as std for normal init in embedding
357+
# FlexAttention does not support dropout (not part of PyTorch's
358+
# flex_attention API). Override to 0 if the model sets it.
359+
if hasattr(self, "attention_dropout") and self.attention_dropout > 0:
360+
logger.warning(
361+
f"attention_dropout={self.attention_dropout} is not supported "
362+
"with FlexAttention, setting to 0.0"
363+
)
364+
self.attention_dropout = 0.0
267365

268366
# When dim is explicitly overridden (e.g. debugmodel), derive the
269367
# dependent sizes from it. Otherwise keep what AutoConfig loaded from
@@ -363,6 +461,15 @@ def __init__(self, config: Config):
363461

364462
for layer in self.model.model.layers.values():
365463
layer.moe_enabled = False
464+
# Attach FlexAttention kernel module to each attention layer
465+
# so apply_cp_to_forward can find and wrap it for CP.
466+
if (
467+
hasattr(layer, "self_attn")
468+
and config.attn_implementation == "flex_torchtitan"
469+
):
470+
layer.self_attn.register_module(
471+
"_flex_kernel", HFFlexAttention(config=HFFlexAttention.Config())
472+
)
366473

367474
def set_cp_mesh(self, mesh):
368475
self.cp_mesh = mesh
@@ -656,16 +763,24 @@ def rotary_emb(self, value):
656763
"Could not find rotary_emb in the model. Please check the model structure."
657764
)
658765

659-
def forward(self, *args, **kwargs):
766+
def forward(self, *args, positions=None, attention_masks=None, **kwargs):
660767
local_seq_len = self.max_seq_len
661768
local_seq_len //= (
662769
self.cp_mesh.size()
663770
if self.cp_mesh is not None and self.cp_mesh.size() > 1
664771
else 1
665772
)
666-
kwargs["position_ids"] = torch.arange(
667-
local_seq_len, device=args[0].device
668-
).unsqueeze(0)
773+
774+
if positions is not None:
775+
kwargs["position_ids"] = positions.to(args[0].device)
776+
else:
777+
kwargs["position_ids"] = torch.arange(
778+
local_seq_len, device=args[0].device
779+
).unsqueeze(0)
780+
781+
if attention_masks is not None:
782+
kwargs["attention_mask"] = attention_masks
783+
669784
output = self.model.model(*args, **kwargs)
670785
if self._skip_lm_head:
671786
return output.last_hidden_state

‎torchtitan/experiments/transformers_modeling_backend/parallelize.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,15 @@ def parallelize_hf_transformers(
101101

102102
if parallel_dims.cp_enabled:
103103
model.set_cp_mesh(parallel_dims.get_mesh("cp"))
104+
# Collect HFFlexAttention kernel modules for CP wrapping
105+
flex_modules = []
106+
for layer in model.layers.values():
107+
if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "_flex_kernel"):
108+
flex_modules.append(layer.self_attn._flex_kernel)
109+
if flex_modules:
110+
from torchtitan.distributed.context_parallel import apply_cp_to_forward
111+
112+
apply_cp_to_forward(flex_modules, parallel_dims.get_mesh("cp"))
104113
logger.info("Applied Context Parallel to the model")
105114

106115
return model

‎torchtitan/experiments/transformers_modeling_backend/tests/integration_tests.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,17 @@ def build_transformers_modeling_backend_test_list() -> list[OverrideDefinitions]
3434
"transformers_modeling_backend_fsdp+tp+pp",
3535
ngpu=8,
3636
),
37+
OverrideDefinitions(
38+
[
39+
[
40+
"--module transformers_modeling_backend",
41+
"--config transformers_modeling_backend_sft_debugmodel",
42+
],
43+
],
44+
"Transformers Backend SFT ChatDataset",
45+
"transformers_modeling_backend_sft",
46+
ngpu=2,
47+
),
3748
]
3849
return integration_tests_flavors
3950

0 commit comments

Comments
 (0)