Skip to content

Commit 6beb16e

Browse files
akoumpaakoumpa
authored
Merge branch 'main' into main
2 parents abe8433 + 174c5d1 commit 6beb16e

64 files changed

Lines changed: 2431 additions & 174 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎examples/vlm_finetune/gemma4/gemma4_4b.yaml‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,11 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
# Configuration for fine-tuning Gemma 4 5B with MedPix dataset for image description
16-
# torchrun --nproc-per-node=2 examples/vlm_finetune/finetune.py -c examples/vlm_finetune/gemma4/gemma4_4b.yaml
15+
# Configuration for fine-tuning Gemma 4 E4B with MedPix dataset for image description
16+
# PR 2237 changed the default reduce_dtype from bfloat16 to fp32, which causes OOMs on E4B with 2 GPUs.
17+
# Although it can be made to run on 2 GPUs by setting reduce_dtype back to bfloat16 and AC=True, its still very
18+
# tight. Hence, using 8 GPUs to run this recipe.
19+
# torchrun --nproc-per-node=8 examples/vlm_finetune/finetune.py -c examples/vlm_finetune/gemma4/gemma4_4b.yaml
1720

1821
recipe: FinetuneRecipeForVLM
1922

‎examples/vlm_finetune/gemma4/gemma4_4b_te.yaml‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,12 @@
1515
# Fine-tuning Gemma4 E4B with TransformerEngine DotProductAttention (FlashAttention-3).
1616
# attn_implementation: te injects TE into all self_attn modules post-init,
1717
# enabling FA3 kernels without any model-specific code changes.
18-
#
19-
# torchrun --nproc-per-node=2 examples/vlm_finetune/finetune.py \
18+
19+
# PR 2237 changed the default reduce_dtype from bfloat16 to fp32, which causes OOMs on E4B with 2 GPUs.
20+
# Although it can be made to run on 2 GPUs by setting reduce_dtype back to bfloat16 and AC=True, its still very
21+
# tight. Hence, using 8 GPUs to run this recipe.
22+
23+
# torchrun --nproc-per-node=8 examples/vlm_finetune/finetune.py \
2024
# -c examples/vlm_finetune/gemma4/gemma4_4b_te.yaml
2125

2226
recipe: FinetuneRecipeForVLM

‎examples/vlm_finetune/gemma4_joint_drafter/gemma4_4b_joint_drafter_medpix.yaml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
# ``_shift_labels_left`` already.
2525
#
2626
# torchrun --nproc-per-node=8 examples/vlm_finetune/finetune.py \
27-
# -c examples/vlm_finetune/gemma4/gemma4_4b_joint_drafter_multi_step.yaml
27+
# -c examples/vlm_finetune/gemma4_joint_drafter/gemma4_4b_joint_drafter_medpix.yaml
2828
#
2929
# Requires: transformers>=5.8.0.dev
3030

@@ -125,7 +125,7 @@ optimizer:
125125
# you see early base-loss spikes when bumping K, bump warmup to 50 first
126126
# before lowering drafter_loss_weight.
127127
lr_scheduler:
128-
lr_warmup_steps: 25
128+
# 10% of max_steps as warmup is used by default which is what we want hence not setting lr_warmup_steps here.
129129
lr_decay_style: cosine
130130

131131
freeze_config:

‎examples/vlm_finetune/gemma4_joint_drafter/gemma4_4b_joint_drafter_tulu_magicoder_mix.yaml‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121
# turns, ``default_collate_fn`` does not produce a ``pixel_values`` tensor and
2222
# the base treats the batch as text-only.
2323
#
24-
# torchrun --nproc-per-node=8 examples/vlm_finetune/finetune.py \
25-
# -c examples/vlm_finetune/gemma4/gemma4_4b_joint_drafter_text_mix.yaml
24+
# torchrun --nproc-per-node=8 examples/vlm_finetune/finetune.py -c \
25+
# examples/vlm_finetune/gemma4_joint_drafter/gemma4_4b_joint_drafter_tulu_magicoder_mix.yaml
2626
#
2727
# Requires: transformers>=5.8.0.dev
2828

@@ -121,9 +121,8 @@ optimizer:
121121
weight_decay: 0.01
122122
betas: [0.9, 0.95]
123123

124-
# 10% of max_steps as warmup
125124
lr_scheduler:
126-
lr_warmup_steps: 200
125+
# 10% of max_steps as warmup is used by default which is what we want hence not setting lr_warmup_steps here.
127126
lr_decay_style: cosine
128127

129128
freeze_config:

‎examples/vlm_finetune/qwen3_5/qwen3_5_9b.yaml‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,11 @@ freeze_config:
7171
freeze_language_model: false
7272

7373
loss_fn:
74-
_target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
74+
# FusedLinearCrossEntropy fuses the lm_head projection with cross-entropy and
75+
# computes it in chunks, so the full [num_tokens, vocab] logits tensor is never
76+
# materialized. MaskedCrossEntropy fp32-upcasts the full logits and OOMs at the
77+
# F.cross_entropy call on large VLM batches (AM-457). Matches qwen3_6_27b_medpix.
78+
_target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy
7579

7680
dataset:
7781
_target_: nemo_automodel.components.datasets.vlm.datasets.make_medpix_dataset

‎examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,10 @@ distributed:
6767
ep_size: 8
6868

6969
sequence_parallel: false
70+
# Recompute transformer-block activations in backward so the fp32 master weights
71+
# (restored for custom MoE under FSDP2 in #1896) fit within 80GB; without this the
72+
# 35B-A3B run OOMs intermittently at its ~67 GiB steady-state peak on 8xH100.
73+
activation_checkpointing: true
7074

7175
freeze_config:
7276
freeze_vision_tower: true

‎examples/vlm_finetune/qwen3_5_moe/qwen3_6_35b.yaml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,10 @@ distributed:
7171
ep_size: 8
7272

7373
sequence_parallel: false
74+
# Recompute transformer-block activations in backward so the fp32 master weights
75+
# (restored for custom MoE under FSDP2 in #1896) fit within 80GB; without this the
76+
# 35B-A3B run OOMs intermittently at its ~67 GiB steady-state peak on 8xH100.
77+
activation_checkpointing: true
7478

7579
freeze_config:
7680
freeze_vision_tower: true

‎nemo_automodel/_transformers/infrastructure.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -551,7 +551,8 @@ def apply_model_infrastructure(
551551
for mp in model_parts:
552552
model_wrapper.maybe_compile(mp)
553553
if isinstance(model_wrapper, DDPManager):
554-
setattr(model.module, "_pre_shard_hf_state_dict_keys", pre_shard_hf_state_dict_keys)
554+
ddp_model = getattr(model, "module", model)
555+
setattr(ddp_model, "_pre_shard_hf_state_dict_keys", pre_shard_hf_state_dict_keys)
555556
else:
556557
setattr(model, "_pre_shard_hf_state_dict_keys", pre_shard_hf_state_dict_keys)
557558

‎nemo_automodel/_transformers/model_init.py‎

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,10 @@ def _check_model_inputs(func):
6060
import nemo_automodel.components.distributed.utils as dist_utils
6161
from nemo_automodel._transformers.registry import ModelRegistry
6262
from nemo_automodel.components.distributed.init_utils import get_local_world_size_preinit, get_world_size_safe
63+
from nemo_automodel.components.models.common.gated_delta_net_fp32 import (
64+
has_gated_delta_net_fp32_checkpoint_contract,
65+
is_gated_delta_net_fp32_param_key,
66+
)
6367
from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
6468
from nemo_automodel.components.utils.model_utils import resolve_trust_remote_code, skip_random_init
6569
from nemo_automodel.shared.utils import dtype_from_str
@@ -662,29 +666,42 @@ def _restore_loaded_model_dtype(
662666
if not checkpoint_dtypes:
663667
return
664668

669+
preserve_gdn_fp32_params = has_gated_delta_net_fp32_checkpoint_contract(hf_config)
665670
restored_dtype_by_tensor_id: dict[int, torch.dtype] = {}
666671
restored_count = 0
667672
for name, checkpoint_dtype in checkpoint_dtypes.items():
668673
tensor = _get_model_tensor(model, name)
669674
if tensor is None:
670675
continue
671676

677+
effective_checkpoint_dtype = (
678+
torch.float32
679+
if checkpoint_dtype.is_floating_point
680+
and preserve_gdn_fp32_params
681+
and is_gated_delta_net_fp32_param_key(name)
682+
else checkpoint_dtype
683+
)
684+
672685
# Record the checkpoint's original dtype on the tensor as the compute-dtype
673686
# hint. Storage may be upcast below (fp32 master weights), which erases the
674687
# dtype HF intended for compute; downstream sharding (fully_shard_by_dtype)
675688
# reads ``_hf_compute_dtype`` to keep intrinsically-fp32 params (e.g. ``A_log``)
676689
# computing in fp32 while the bulk computes in mp_policy.param_dtype.
677-
if checkpoint_dtype.is_floating_point and tensor.dtype.is_floating_point:
678-
tensor._hf_compute_dtype = checkpoint_dtype
690+
if effective_checkpoint_dtype.is_floating_point and tensor.dtype.is_floating_point:
691+
tensor._hf_compute_dtype = effective_checkpoint_dtype
679692

680693
# Pick the unification target. For an explicit floating request, take the
681694
# wider of (checkpoint, requested) so explicit fp32 is honored as master
682695
# weights while intrinsically-fp32 checkpoint params survive a bf16 request.
683696
# For "auto" (requested_dtype is None) or non-floating tensors, mirror the
684697
# checkpoint dtype exactly (preserves today's behavior).
685-
target_dtype = checkpoint_dtype
686-
if requested_dtype is not None and checkpoint_dtype.is_floating_point and tensor.dtype.is_floating_point:
687-
target_dtype = torch.promote_types(checkpoint_dtype, requested_dtype)
698+
target_dtype = effective_checkpoint_dtype
699+
if (
700+
requested_dtype is not None
701+
and effective_checkpoint_dtype.is_floating_point
702+
and tensor.dtype.is_floating_point
703+
):
704+
target_dtype = torch.promote_types(effective_checkpoint_dtype, requested_dtype)
688705

689706
if tensor.dtype == target_dtype:
690707
continue

‎nemo_automodel/components/_peft/lora.py‎

Lines changed: 32 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -300,12 +300,10 @@ def forward(self, x):
300300
use_memory_efficient_lora = self._should_use_memory_efficient_lora(x)
301301
if use_memory_efficient_lora:
302302
if self.dropout_position == "pre" or not self.training or self.dropout_p == 0.0:
303-
return LoRATritonFunction.apply(
304-
x, self.lora_A.weight, self.lora_B.weight, self.scale, x.dtype, False, res
303+
return apply_memory_efficient_lora(
304+
x, self.lora_A.weight, self.lora_B.weight, self.scale, False, res
305305
)
306-
lora_res = LoRATritonFunction.apply(
307-
x, self.lora_A.weight, self.lora_B.weight, self.scale, x.dtype, False
308-
)
306+
lora_res = apply_memory_efficient_lora(x, self.lora_A.weight, self.lora_B.weight, self.scale, False)
309307
else:
310308
lora_res = self.lora_B(self.lora_A(x) * self.scale)
311309
if self.dropout_position == "post":
@@ -393,10 +391,8 @@ def forward(self, x):
393391
x = F.dropout(x, p=self.dropout_p, training=self.training)
394392
if self.use_memory_efficient_lora:
395393
if self.dropout_position == "pre" or not self.training or self.dropout_p == 0.0:
396-
return LoRATritonFunction.apply(
397-
x, self.lora_A.weight, self.lora_B.weight, self.scale, x.dtype, True, res
398-
)
399-
lora_res = LoRATritonFunction.apply(x, self.lora_A.weight, self.lora_B.weight, self.scale, x.dtype, True)
394+
return apply_memory_efficient_lora(x, self.lora_A.weight, self.lora_B.weight, self.scale, True, res)
395+
lora_res = apply_memory_efficient_lora(x, self.lora_A.weight, self.lora_B.weight, self.scale, True)
400396
else:
401397
lora_res = self.lora_B(self.lora_A(x) * self.scale)
402398
if self.dropout_position == "post":
@@ -704,11 +700,14 @@ def forward(x, lora_A, lora_B, scale, dtype, use_triton_kernel=True, res=None):
704700
705701
Reshapes 3D tensors into 2D and then calls either Triton kernels or PyTorch matmuls. When ``res`` is
706702
provided, the residual is added in-place into the LoRA output to avoid allocating a separate add result.
703+
704+
Always returns a **2D** tensor; the caller restores the original leading dimensions. Keeping the
705+
``(N, out) -> (bs, seq, out)`` reshape *outside* this ``autograd.Function`` means the Function's output
706+
is never a view, so a downstream consumer may safely mutate the LoRA output in place (the reshape done
707+
by the caller is an ordinary autograd view, which supports in-place ops).
707708
"""
708-
reshape = x.dim() == 3
709-
if reshape:
710-
bs, seq_len, d = x.shape
711-
x = x.reshape(-1, d)
709+
if x.dim() == 3:
710+
x = x.reshape(-1, x.shape[-1])
712711
if res is not None:
713712
res = res.reshape(-1, res.shape[-1])
714713

@@ -720,8 +719,6 @@ def forward(x, lora_A, lora_B, scale, dtype, use_triton_kernel=True, res=None):
720719
if res is not None:
721720
lora_res.add_(res)
722721

723-
if reshape:
724-
return lora_res.view(bs, seq_len, -1)
725722
return lora_res
726723

727724
@staticmethod
@@ -734,11 +731,16 @@ def backward(ctx, d_y):
734731
"""
735732
x, lora_A, lora_B = ctx.saved_tensors
736733
scale = ctx.scale
737-
d_res = d_y if ctx.has_residual and ctx.needs_input_grad[6] else None
738734

739735
reshape = x.dim() == 3
740736
if reshape:
741737
bs, seq_len, d = x.shape
738+
# forward now returns a 2D output, so d_y arrives 2D; the residual input kept its original
739+
# (possibly 3D) shape, so its gradient must be reshaped back to match that input.
740+
d_res = None
741+
if ctx.has_residual and ctx.needs_input_grad[6]:
742+
d_res = d_y.reshape(bs, seq_len, -1) if reshape else d_y
743+
if reshape:
742744
d_y = d_y.reshape(-1, d_y.shape[-1])
743745
x = x.reshape(-1, d)
744746

@@ -770,3 +772,17 @@ def backward(ctx, d_y):
770772
if ctx.num_inputs == 6:
771773
return gradients + (None,)
772774
return gradients
775+
776+
777+
def apply_memory_efficient_lora(x, lora_A, lora_B, scale, use_triton_kernel, res=None):
778+
"""Run :class:`LoRATritonFunction` and restore the input's leading dimensions.
779+
780+
``LoRATritonFunction.forward`` returns a 2D tensor (its reshape is intentionally kept outside the
781+
autograd Function so the output is never a view). Reshape back to the input rank here; the result
782+
is an ordinary autograd view, which — unlike a custom-Function output view — a downstream consumer
783+
may mutate in place (e.g. transformers' gemma3n ``project_per_layer_inputs``).
784+
"""
785+
out = LoRATritonFunction.apply(x, lora_A, lora_B, scale, x.dtype, use_triton_kernel, res)
786+
if x.dim() == 3:
787+
out = out.reshape(*x.shape[:-1], -1)
788+
return out

0 commit comments

Comments
 (0)