Skip to content

Commit a4f945f

Browse files
qgallouedecqgallouedec
authored
Add Beam: tiny model, tests and SFT example (#7629)
1 parent c803af7 commit a4f945f

14 files changed

Lines changed: 256 additions & 3 deletions

‎README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222

2323
## 🎉 What's New
2424

25+
**🔥 Beam:** Fine-tune [Beam](https://huggingface.co/reflection/Beam-1.0-501B-A23B) with [`SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer) on a single node of 8 B300s. See the [LoRA example](https://github.com/huggingface/trl/tree/main/examples/sft_beam).
26+
2527
**🔥 Muse Spark Open:** Fine-tune [Muse Spark Open](https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B) with [`SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer), sharded across your GPUs at load time. See the [8-node LoRA example](https://github.com/huggingface/trl/tree/main/examples/sft_muse_spark_open).
2628

2729
**🚀 Training on AWS Trainium:** TRL's trainers now run on [AWS Trainium](https://huggingface.co/docs/trl/aws_trainium) through TorchNeuron, PyTorch's native Neuron backend.

‎docs/source/example_overview.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ Check for additional optional dependencies [here](https://github.com/huggingface
5454
| [`rloo_visual_math`](https://github.com/huggingface/trl/tree/main/examples/rloo_visual_math) | RLOO fine-tuning of a multimodal model for reasoning using the [lmms-lab/multimodal-open-r1-8k-verified](https://huggingface.co/datasets/lmms-lab/multimodal-open-r1-8k-verified) dataset. | |
5555
| [`sdft_privileged_context`](https://github.com/huggingface/trl/tree/main/examples/sdft_privileged_context) | Self-distillation fine-tuning with [`experimental.sdft.SDFTTrainer`], distilling privileged (teacher-only) context into the model. | |
5656
| [`sdpo_math`](https://github.com/huggingface/trl/tree/main/examples/sdpo_math) | SDPO with [`experimental.sdpo.SDPOTrainer`] using verifiable math rewards and optional environment feedback on [openai/gsm8k](https://huggingface.co/datasets/openai/gsm8k). | |
57+
| [`sft_beam`](https://github.com/huggingface/trl/tree/main/examples/sft_beam) | LoRA SFT of reflection/Beam-1.0-501B-A23B on one 8xB300 node with [`SFTTrainer`], sharded at load time by transformers' `DistributedConfig` (expert parallelism with token dispatch). | |
5758
| [`sft_diffusion_gemma`](https://github.com/huggingface/trl/tree/main/examples/sft_diffusion_gemma) | SFT of the DiffusionGemma block-diffusion language model on GSM8K by extending the [`SFTTrainer`] with a block-diffusion objective. | |
5859
| [`sft_gemma3`](https://github.com/huggingface/trl/tree/main/examples/sft_gemma3) | SFT of Gemma 3 on the Codeforces COTS dataset. | |
5960
| [`sft_gemma3_vision`](https://github.com/huggingface/trl/tree/main/examples/sft_gemma3_vision) | SFT of Gemma 3 on vision to text tasks. | |

‎docs/source/grpo_trainer.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -961,6 +961,7 @@ The returned images are automatically injected into the conversation and passed
961961

962962
Tested with:
963963

964+
- **Beam** — e.g., `reflection/Beam-1.0-501B-A23B`
964965
- [**Gemma4**](https://huggingface.co/collections/google/gemma-4) — e.g., `google/gemma-4-E2B-it`
965966
- **GLM-4-MoE** ([4.5](https://huggingface.co/collections/zai-org/glm-45), [4.6](https://huggingface.co/collections/zai-org/glm-46) or [4.7](https://huggingface.co/collections/zai-org/glm-47)) — e.g., `zai-org/GLM-4.7`
966967
- [**GPT-OSS**](https://huggingface.co/collections/openai/gpt-oss) — e.g., `openai/gpt-oss-20b`

‎docs/source/index.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ The library is integrated with 🤗 [transformers](https://github.com/huggingfac
1212

1313
## 🎉 What's New
1414

15+
**🔥 Beam:** Fine-tune [Beam](https://huggingface.co/reflection/Beam-1.0-501B-A23B) with [`SFTTrainer`] on a single node of 8 B300s. See the [LoRA example](https://github.com/huggingface/trl/tree/main/examples/sft_beam).
16+
1517
**🔥 Muse Spark Open:** Fine-tune [Muse Spark Open](https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B) with [`SFTTrainer`], sharded across your GPUs at load time. See the [8-node LoRA example](https://github.com/huggingface/trl/tree/main/examples/sft_muse_spark_open).
1618

1719
**🚀 Training on AWS Trainium:** TRL's trainers now run on [AWS Trainium](aws_trainium) through TorchNeuron, PyTorch's native Neuron backend.

‎examples/sft_beam/sft_beam.py‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# /// script
16+
# dependencies = [
17+
# "trl>=1.16.0",
18+
# "transformers>=5.20.1",
19+
# "peft>=0.21.0",
20+
# "accelerate>=1.15.0",
21+
# "flash-attn-4",
22+
# "trackio",
23+
# ]
24+
# ///
25+
26+
# docstyle-ignore
27+
"""
28+
LoRA SFT of Beam on one 8xB300 node: the experts are split 8 ways, everything else is sharded across the 8 GPUs,
29+
and each rank trains on its own slice of the batch. One epoch is 36 steps of 1.75 s, 32.8k tokens per step,
30+
223 GiB per GPU.
31+
32+
sbatch sft_beam.slurm
33+
"""
34+
35+
from datasets import load_dataset
36+
from peft import LoraConfig
37+
from transformers.distributed import DistributedConfig
38+
39+
from trl import SFTConfig, SFTTrainer
40+
41+
42+
MODEL = "reflection/Beam-1.0-501B-A23B"
43+
44+
# Assistant turns carry their reasoning in a `thinking` field, which the chat template renders into the analysis
45+
# channel.
46+
train_dataset = load_dataset("HuggingFaceH4/Multilingual-Thinking", split="train")
47+
48+
training_args = SFTConfig(
49+
output_dir="Beam-SFT",
50+
model_init_kwargs={
51+
# TRL loads a model id in float32 unless told otherwise; "auto" keeps the checkpoint's dtype.
52+
"dtype": "auto",
53+
# Flash Attention 4 runs on Blackwell and keeps packed sequences apart.
54+
"attn_implementation": "flash_attention_4",
55+
# The model's own expert plan selects token dispatch, so the sizes are the whole configuration.
56+
"distributed_config": DistributedConfig(fsdp_size=8, ep_size=8),
57+
},
58+
per_device_train_batch_size=1,
59+
# B300s have the memory to keep the activations: 1.4x faster than recomputing them.
60+
gradient_checkpointing=False,
61+
num_train_epochs=1,
62+
max_length=4096,
63+
# bfd packing flattens the batch (padding-free) and passes position_ids, which the attention uses to keep the
64+
# packed sequences apart.
65+
packing=True,
66+
# Train on the assistant turns only (reasoning, answer and the end-of-turn token).
67+
assistant_only_loss=True,
68+
learning_rate=1e-4,
69+
logging_steps=1,
70+
# Resuming is not supported for models sharded at load time, so the optimizer state cannot be checkpointed.
71+
save_only_model=True,
72+
report_to="trackio",
73+
)
74+
trainer = SFTTrainer(
75+
model=MODEL,
76+
args=training_args,
77+
train_dataset=train_dataset,
78+
peft_config=LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]),
79+
)
80+
trainer.train()

‎examples/sft_beam/sft_beam.slurm‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#!/bin/bash
2+
# Submit from this directory: sbatch sft_beam.slurm
3+
#SBATCH --nodes=1
4+
#SBATCH --gres=gpu:8
5+
#SBATCH --ntasks-per-node=1
6+
#SBATCH --mem=0
7+
#SBATCH --time=02:00:00
8+
9+
torchrun --nproc_per_node 8 "$SLURM_SUBMIT_DIR/sft_beam.py"

‎scripts/generate_tiny_models/for_causal_lm/beam_for_causal_lm.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import torch
16+
from transformers import AutoConfig, AutoTokenizer, GenerationConfig
17+
from transformers.models.beam import BeamConfig, BeamForCausalLM
18+
19+
from .._common import (
20+
check_dtype_pattern,
21+
check_transformers_version,
22+
init_weights_tiny_model,
23+
print_config_diff,
24+
push_to_hub,
25+
set_seed,
26+
smoke_test,
27+
)
28+
29+
30+
check_transformers_version("5.20.1")
31+
32+
set_seed()
33+
34+
MODEL_ID = "reflection/Beam-1.0-501B-A23B"
35+
36+
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
37+
generation_config = GenerationConfig.from_pretrained(MODEL_ID)
38+
reference = AutoConfig.from_pretrained(MODEL_ID)
39+
config = BeamConfig(
40+
vocab_size=reference.vocab_size,
41+
hidden_size=16,
42+
intermediate_size=32,
43+
moe_intermediate_size=8,
44+
num_hidden_layers=4,
45+
num_attention_heads=4,
46+
num_key_value_heads=2,
47+
head_dim=16,
48+
num_experts=8,
49+
num_experts_per_tok=2,
50+
# Group-limited routing needs at least 2 experts per group (the group score sums the top 2)
51+
n_group=2,
52+
topk_group=1,
53+
sliding_window=4,
54+
# One dense layer then sparse ones, and both attention types, with a sliding window shorter than the test inputs
55+
mlp_layer_types=["dense", "sparse", "sparse", "sparse"],
56+
layer_types=["sliding_attention", "sliding_attention", "sliding_attention", "full_attention"],
57+
rope_parameters=reference.rope_parameters,
58+
final_logit_softcapping=reference.final_logit_softcapping,
59+
routed_scaling_factor=reference.routed_scaling_factor,
60+
max_position_embeddings=reference.max_position_embeddings,
61+
tie_word_embeddings=reference.tie_word_embeddings,
62+
)
63+
model = BeamForCausalLM(config).to(dtype=torch.bfloat16)
64+
init_weights_tiny_model(model)
65+
# The routers stay in float32, as in the reference checkpoint
66+
for name, module in model.named_modules():
67+
if name.endswith("mlp.gate"):
68+
module.float()
69+
smoke_test(model, tokenizer)
70+
check_dtype_pattern(MODEL_ID, model)
71+
print_config_diff(MODEL_ID, model)
72+
push_to_hub(model, tokenizer, generation_config, "tiny")

‎tests/test_chat_template_utils.py‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -393,6 +393,14 @@ class TestSupportsToolCalling:
393393
@pytest.mark.parametrize(
394394
"model_id",
395395
[
396+
pytest.param(
397+
"trl-internal-testing/tiny-BeamForCausalLM",
398+
id="beam",
399+
marks=pytest.mark.skipif(
400+
Version(transformers.__version__) < Version("5.20.1"),
401+
reason="Beam was introduced in transformers-5.20.1",
402+
),
403+
),
396404
pytest.param("trl-internal-testing/tiny-DeepseekV3ForCausalLM", id="deepseekv3"),
397405
pytest.param("trl-internal-testing/tiny-DeepseekV3ForCausalLM-0528", id="deepseekv3-0528"),
398406
pytest.param(
@@ -806,6 +814,14 @@ def test_template_error_returns_false(self):
806814
@pytest.mark.parametrize(
807815
"tokenizer_name",
808816
[
817+
pytest.param(
818+
"trl-internal-testing/tiny-BeamForCausalLM",
819+
id="beam",
820+
marks=pytest.mark.skipif(
821+
Version(transformers.__version__) < Version("5.20.1"),
822+
reason="Beam was introduced in transformers-5.20.1",
823+
),
824+
),
809825
pytest.param("trl-internal-testing/tiny-CohereForCausalLM", id="cohere"),
810826
# Real Cohere checkpoints ship their chat template as named variants, which load as a dict
811827
pytest.param(
@@ -1292,6 +1308,14 @@ def test_unsupported_default_variant(self):
12921308
@pytest.mark.parametrize(
12931309
"model_name",
12941310
[
1311+
pytest.param(
1312+
"trl-internal-testing/tiny-BeamForCausalLM",
1313+
id="beam",
1314+
marks=pytest.mark.skipif(
1315+
Version(transformers.__version__) < Version("5.20.1"),
1316+
reason="Beam was introduced in transformers-5.20.1",
1317+
),
1318+
),
12951319
pytest.param("trl-internal-testing/tiny-Glm4MoeForCausalLM", id="glm4moe"),
12961320
pytest.param("trl-internal-testing/tiny-GptOssForCausalLM", id="gptoss"),
12971321
pytest.param(
@@ -1476,12 +1500,13 @@ def test_parse_response_with_reasoning_content(self, model_name):
14761500
pytest.skip("gpt-oss thinking/content separation requires the response_template parser (>= 5.13).")
14771501

14781502
processing_class = self._load(model_name)
1479-
# gpt-oss and LFM2.5 use the `thinking` field name (matching their own chat templates, which read it back off
1480-
# the message) rather than the `reasoning_content` convention used by other models.
1503+
# gpt-oss, Beam and LFM2.5 use the `thinking` field name (matching their own chat templates, which read it back
1504+
# off the message) rather than the `reasoning_content` convention used by other models.
14811505
reasoning_field = (
14821506
"thinking"
14831507
if model_name
14841508
in (
1509+
"trl-internal-testing/tiny-BeamForCausalLM",
14851510
"trl-internal-testing/tiny-GptOssForCausalLM",
14861511
"trl-internal-testing/tiny-Lfm2ForCausalLM-2.5",
14871512
"trl-internal-testing/tiny-Lfm2VlForConditionalGeneration-2.5",
@@ -1554,7 +1579,10 @@ def test_parse_response_tool_call_with_content(self, model_name):
15541579
processing_class = self._load(model_name)
15551580
tool_calls = [{"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}}]
15561581
assistant = {"role": "assistant", "tool_calls": tool_calls}
1557-
if model_name == "trl-internal-testing/tiny-GptOssForCausalLM":
1582+
if model_name in (
1583+
"trl-internal-testing/tiny-BeamForCausalLM",
1584+
"trl-internal-testing/tiny-GptOssForCausalLM",
1585+
):
15581586
assistant["thinking"] = "Let's call the tool."
15591587
assistant["content"] = ""
15601588
elif model_name == "trl-internal-testing/tiny-MuseGlimmerForConditionalGeneration":

‎tests/test_data_utils.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -428,6 +428,13 @@ def test_negative_2(self):
428428

429429
class TestApplyChatTemplate(TrlTestCase):
430430
tokenizers = [
431+
pytest.param(
432+
"trl-internal-testing/tiny-BeamForCausalLM",
433+
marks=pytest.mark.skipif(
434+
Version(transformers.__version__) < Version("5.20.1"),
435+
reason="Beam was introduced in transformers-5.20.1",
436+
),
437+
),
431438
"trl-internal-testing/tiny-CohereForCausalLM",
432439
"trl-internal-testing/tiny-Cohere2ForCausalLM",
433440
"trl-internal-testing/tiny-DeepseekV3ForCausalLM",

‎tests/test_dpo_trainer.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,13 @@ class TestDPOTrainer(TrlTestCase):
183183
"trl-internal-testing/tiny-GptOssForCausalLM",
184184
"trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
185185
"trl-internal-testing/tiny-Qwen3MoeForCausalLM",
186+
pytest.param(
187+
"trl-internal-testing/tiny-BeamForCausalLM",
188+
marks=pytest.mark.skipif(
189+
Version(transformers.__version__) < Version("5.20.1"),
190+
reason="Beam was introduced in transformers-5.20.1",
191+
),
192+
),
186193
pytest.param(
187194
"trl-internal-testing/tiny-NemotronHForCausalLM-nano",
188195
marks=pytest.mark.skipif(

0 commit comments

Comments
 (0)