Skip to content

Commit 09762b6

Browse files
zyzhou5zyzhou5
committed
feat(bagel): add multimodal training support
Add BAGEL model, dataset, recipe, examples, EMA support, and focused unit coverage as a single public-facing change. The example token budgets match the one-node AM parity configuration to avoid the current full-budget OOM path. Validation: Native BAGEL and AM PR-ready 3k-step one-node parity runs completed successfully in W&B project Bagel-One-Node-PR-Parity.
1 parent 14f14cd commit 09762b6

40 files changed

Lines changed: 8329 additions & 3 deletions

‎examples/multimodal_finetune/bagel/bagel_sft.yaml‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# BAGEL fine-tuning - joint understanding + generation.
4+
#
5+
# This is the default AM BAGEL fine-tuning recipe. It activates the MoT decoder,
6+
# flow-matching MSE loss, EMA, and the VAE encode path for joint understanding
7+
# and generation training.
8+
9+
recipe: FinetuneRecipeForMultimodal
10+
11+
seed: 4396
12+
13+
step_scheduler:
14+
global_batch_size: 8
15+
local_batch_size: 1
16+
max_steps: 100
17+
ckpt_every_steps: 100000
18+
val_every_steps: 100000
19+
save_checkpoint_every_epoch: false
20+
21+
dist_env:
22+
backend: nccl
23+
timeout_minutes: 30
24+
25+
model:
26+
_target_: nemo_automodel.NeMoAutoModelForMultimodalLM.from_pretrained
27+
pretrained_model_name_or_path: ByteDance-Seed/BAGEL-7B-MoT
28+
torch_dtype: float32
29+
stage: 2
30+
# Optional override. If omitted, the recipe uses ae.safetensors from the
31+
# local checkpoint directory or downloads it from the same HF repo.
32+
# vae_path: /path/to/ae.safetensors
33+
# max_latent_size: 64
34+
35+
checkpoint:
36+
enabled: false
37+
checkpoint_dir: checkpoints/bagel-sft-am
38+
model_save_format: torch_save
39+
40+
distributed:
41+
strategy: fsdp2
42+
dp_size: none
43+
dp_replicate_size: 1
44+
cp_size: 1
45+
sequence_parallel: false
46+
activation_checkpointing: true
47+
mp_policy:
48+
param_dtype: bfloat16
49+
reduce_dtype: bfloat16
50+
output_dtype: bfloat16
51+
cast_forward_inputs: true
52+
53+
dataset:
54+
# Replace these paths with prepared BAGEL-format data.
55+
dataset_info:
56+
t2i_pretrain:
57+
t2i:
58+
data_dir: /path/to/bagel_data/t2i
59+
num_files: 8
60+
num_total_samples: 1000
61+
unified_edit:
62+
edit_data:
63+
data_dir: /path/to/bagel_data/editing/images
64+
num_files: 8
65+
num_total_samples: 1000
66+
parquet_info_path: /path/to/bagel_data/editing/parquet_info.json
67+
vlm_sft:
68+
vlm_data:
69+
data_dir: /path/to/bagel_data/vlm/images
70+
jsonl_path: /path/to/bagel_data/vlm/train.jsonl
71+
num_total_samples: 1000
72+
grouped_datasets:
73+
t2i_pretrain:
74+
dataset_names:
75+
- t2i
76+
image_transform_args:
77+
image_stride: 16
78+
max_image_size: 1024
79+
min_image_size: 512
80+
is_mandatory: true
81+
num_used_data:
82+
- 8
83+
weight: 1
84+
unified_edit:
85+
dataset_names:
86+
- edit_data
87+
image_transform_args:
88+
image_stride: 16
89+
max_image_size: 1024
90+
min_image_size: 512
91+
vit_image_transform_args:
92+
image_stride: 14
93+
max_image_size: 518
94+
min_image_size: 224
95+
is_mandatory: false
96+
num_used_data:
97+
- 8
98+
weight: 1
99+
vlm_sft:
100+
dataset_names:
101+
- vlm_data
102+
image_transform_args:
103+
image_stride: 14
104+
max_image_size: 980
105+
min_image_size: 378
106+
max_pixels: 2_007_040
107+
frame_sampler_args:
108+
max_num_frames: 12
109+
min_num_frames: 8
110+
is_mandatory: true
111+
shuffle_lines: true
112+
shuffle_seed: 0
113+
num_used_data:
114+
- 1000
115+
weight: 1
116+
expected_num_tokens: 10240
117+
max_num_tokens: 11520
118+
max_num_tokens_per_sample: 10240
119+
prefer_buffer_before: 16384
120+
max_buffer_size: 50
121+
vit_patch_size: 14
122+
max_num_patch_per_side: 70
123+
interpolate_pos: false
124+
data_seed: 42
125+
text_cond_dropout_prob: 0.1
126+
vit_cond_dropout_prob: 0.4
127+
vae_cond_dropout_prob: 0.1
128+
# VAE-side downsample: vae.downsample (8) * latent_patch_size (2) = 16.
129+
vae_image_downsample: 16
130+
max_latent_size: 64
131+
use_flex: true
132+
133+
dataloader:
134+
num_workers: 1
135+
pin_memory: true
136+
prefetch_factor: 2
137+
138+
optimizer:
139+
_target_: torch.optim.AdamW
140+
lr: 2.0e-5
141+
betas: [0.9, 0.95]
142+
eps: 1.0e-15
143+
weight_decay: 0.0
144+
145+
lr_scheduler:
146+
schedule: constant
147+
warmup_steps: 2000
148+
149+
clip_grad_norm:
150+
max_norm: 1.0
151+
152+
freeze_config:
153+
freeze_vision_tower: false
154+
freeze_language_model: false
155+
156+
# Exponential moving average of model parameters. The shadow tensors are updated
157+
# in-place after every optimizer.step. Shadow memory is approximately model
158+
# memory, sharded across DP ranks.
159+
# Remove the section to disable EMA tracking.
160+
ema:
161+
decay: 0.9999
162+
163+
wandb:
164+
project: bagel-finetuning
165+
mode: disabled
166+
name: bagel_7b_mot_sft

‎examples/multimodal_finetune/finetune.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
# Copyright (c) 2025, NVIDIA CORPORATION. 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+
"""Thin entry point for multimodal fine-tuning.
16+
17+
Recommended invocation (uses AM's CLI launcher):
18+
19+
automodel examples/multimodal_finetune/bagel/bagel_sft.yaml --nproc-per-node 8
20+
21+
This script remains available as a direct entry point:
22+
23+
python -m torch.distributed.run --nproc_per_node=8 \\
24+
examples/multimodal_finetune/finetune.py \\
25+
-c examples/multimodal_finetune/bagel/bagel_sft.yaml
26+
"""
27+
28+
from __future__ import annotations
29+
30+
import warnings
31+
32+
warnings.warn(
33+
"Running recipes via examples/ scripts is deprecated. "
34+
"Use: automodel <config.yaml> [--nproc-per-node N]\n"
35+
"See BREAKING_CHANGES.md for details.",
36+
DeprecationWarning,
37+
stacklevel=2,
38+
)
39+
40+
from nemo_automodel.components.config._arg_parser import parse_args_and_load_config
41+
from nemo_automodel.recipes.multimodal.finetune import FinetuneRecipeForMultimodal
42+
43+
44+
def main(config: str = "examples/multimodal_finetune/bagel/bagel_sft.yaml") -> None:
45+
cfg = parse_args_and_load_config(config)
46+
recipe = FinetuneRecipeForMultimodal(cfg)
47+
recipe.setup()
48+
recipe.run_train_validation_loop()
49+
50+
51+
if __name__ == "__main__":
52+
main()

‎examples/multimodal_pretrain/bagel/bagel_pretrain.yaml‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# BAGEL pretraining - joint understanding + generation.
4+
#
5+
# This config initializes BAGEL from pretrained Qwen and SigLIP backbones,
6+
# then trains on the packed multimodal mixture defined below.
7+
8+
recipe: PretrainRecipeForMultimodal
9+
10+
seed: 4396
11+
12+
step_scheduler:
13+
global_batch_size: 8
14+
local_batch_size: 1
15+
max_steps: 100
16+
ckpt_every_steps: 100000
17+
val_every_steps: 100000
18+
save_checkpoint_every_epoch: false
19+
20+
dist_env:
21+
backend: nccl
22+
timeout_minutes: 30
23+
24+
tokenizer:
25+
pretrained_model_name_or_path: Qwen/Qwen2.5-7B-Instruct
26+
27+
model:
28+
init_mode: hf_backbones
29+
stage: 2
30+
llm_path: Qwen/Qwen2.5-7B-Instruct
31+
vit_path: HuggingFaceM4/siglip-so400m-14-980-flash-attn2-navit
32+
vae_path: /path/to/ae.safetensors
33+
layer_module: Qwen2MoTDecoderLayer
34+
llm_qk_norm: true
35+
tie_word_embeddings: false
36+
copy_init_moe: true
37+
torch_dtype: float32
38+
max_latent_size: 64
39+
40+
checkpoint:
41+
enabled: false
42+
checkpoint_dir: checkpoints/bagel-pretrain-am
43+
model_save_format: torch_save
44+
45+
distributed:
46+
strategy: fsdp2
47+
dp_size: none
48+
dp_replicate_size: 1
49+
cp_size: 1
50+
sequence_parallel: false
51+
activation_checkpointing: true
52+
mp_policy:
53+
param_dtype: bfloat16
54+
reduce_dtype: bfloat16
55+
output_dtype: bfloat16
56+
cast_forward_inputs: true
57+
58+
dataset:
59+
# Replace these paths with prepared BAGEL-format data.
60+
# For parquet-backed mandatory groups, num_files / num_used_data should be
61+
# at least the data-parallel world size so every rank receives files.
62+
dataset_info:
63+
t2i_pretrain:
64+
t2i:
65+
data_dir: /path/to/bagel_data/t2i
66+
num_files: 8
67+
num_total_samples: 1000
68+
unified_edit:
69+
edit_data:
70+
data_dir: /path/to/bagel_data/editing/images
71+
num_files: 8
72+
num_total_samples: 1000
73+
parquet_info_path: /path/to/bagel_data/editing/parquet_info.json
74+
vlm_sft:
75+
vlm_data:
76+
data_dir: /path/to/bagel_data/vlm/images
77+
jsonl_path: /path/to/bagel_data/vlm/train.jsonl
78+
num_total_samples: 1000
79+
grouped_datasets:
80+
t2i_pretrain:
81+
dataset_names:
82+
- t2i
83+
image_transform_args:
84+
image_stride: 16
85+
max_image_size: 1024
86+
min_image_size: 512
87+
is_mandatory: true
88+
num_used_data:
89+
- 8
90+
weight: 1
91+
unified_edit:
92+
dataset_names:
93+
- edit_data
94+
image_transform_args:
95+
image_stride: 16
96+
max_image_size: 1024
97+
min_image_size: 512
98+
vit_image_transform_args:
99+
image_stride: 14
100+
max_image_size: 518
101+
min_image_size: 224
102+
is_mandatory: false
103+
num_used_data:
104+
- 8
105+
weight: 1
106+
vlm_sft:
107+
dataset_names:
108+
- vlm_data
109+
image_transform_args:
110+
image_stride: 14
111+
max_image_size: 980
112+
min_image_size: 378
113+
max_pixels: 2_007_040
114+
frame_sampler_args:
115+
max_num_frames: 12
116+
min_num_frames: 8
117+
is_mandatory: true
118+
shuffle_lines: true
119+
shuffle_seed: 0
120+
num_used_data:
121+
- 1000
122+
weight: 1
123+
expected_num_tokens: 10240
124+
max_num_tokens: 11520
125+
max_num_tokens_per_sample: 10240
126+
prefer_buffer_before: 16384
127+
max_buffer_size: 50
128+
vit_patch_size: 14
129+
max_num_patch_per_side: 70
130+
interpolate_pos: false
131+
data_seed: 42
132+
text_cond_dropout_prob: 0.1
133+
vit_cond_dropout_prob: 0.4
134+
vae_cond_dropout_prob: 0.1
135+
vae_image_downsample: 16
136+
max_latent_size: 64
137+
use_flex: true
138+
139+
dataloader:
140+
num_workers: 1
141+
pin_memory: true
142+
prefetch_factor: 2
143+
144+
optimizer:
145+
_target_: torch.optim.AdamW
146+
lr: 1.0e-4
147+
betas: [0.9, 0.95]
148+
eps: 1.0e-15
149+
weight_decay: 0.0
150+
151+
lr_scheduler:
152+
schedule: constant
153+
warmup_steps: 2000
154+
155+
clip_grad_norm:
156+
max_norm: 1.0
157+
158+
freeze_config:
159+
freeze_vision_tower: false
160+
freeze_language_model: false
161+
162+
ema:
163+
decay: 0.9999
164+
165+
wandb:
166+
project: bagel-pretraining
167+
mode: disabled
168+
name: bagel_mot_pretrain

0 commit comments

Comments
 (0)