fix(vlm): support packed THD Step3.7 pipeline training by new-TonyWang · Pull Request #3572 · NVIDIA-NeMo/Automodel

new-TonyWang · 2026-08-18T08:08:59Z

Summary

add robust text-only THD packing support for Step3.7 VLM SFT, including malformed-row fallback and packed metadata propagation
make pipeline schedules preserve THD auxiliary inputs and scalar losses across pipeline stages
harden Step3.5/Step3.7 MoE paths for FSDP CPU offload and grouped expert execution
improve rank-aware metric/TensorBoard logging and distributed gradient utilities
extend Dion/Muon configuration for the low-memory stateless path
document the Step3.7 text THD packing workflow

Validation

focused unit suite covering datasets, VLM packing, pipeline helpers, THD utilities, Step3.5/Step3.7 model behavior, and optimizer configuration
9 targeted test modules pass
end-to-end 8-GPU Step3.7 packed THD SFT completed 1,000 optimizer steps with PP4/EP2

Notes

The end-to-end validation used mu=0 for the stateless Muon path. This PR does not claim equivalent memory behavior for momentum Muon configurations.

copy-pr-bot · 2026-08-18T08:09:03Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

HuiyingLi · 2026-08-18T16:58:26Z

Thank you @new-TonyWang , could you please resolve the merge conflicts? Thank you!

jgerh

Completed tech pubs review of docs/STEP3P7_TEXT_THD_PACKING.md. Provided a few copyedits and formatting updates to better align with our style guide.

jgerh · 2026-08-18T23:13:48Z

@@ -0,0 +1,63 @@
+# Step-3.7 text SFT with THD packed sequences


Note: The file naming convention for Markdown files is lowercase with hyphens (kebab-case). Please rename STEP3P7_TEXT_THD_PACKING.md to step3p7-text-thd-packing.md, and update any cross-references or index.md entries that point to it.

Suggested change

# Step-3.7 text SFT with THD packed sequences

# Step-3.7 Text SFT with THD Packed Sequences

jgerh · 2026-08-18T23:14:11Z

@@ -0,0 +1,63 @@
+# Step-3.7 text SFT with THD packed sequences
+
+AutoModel's VLM packing pipeline can now consume tokenizer-aware datasets such


Suggested change

AutoModel's VLM packing pipeline can now consume tokenizer-aware datasets such

NeMo AutoModel's VLM packing pipeline can now consume tokenizer-aware datasets such

jgerh · 2026-08-18T23:14:20Z

+AutoModel's VLM packing pipeline can now consume tokenizer-aware datasets such
+as `ChatDataset` without re-tokenizing them through the multimodal processor.
+
+## Data flow


Suggested change

## Data flow

## Data Flow

jgerh · 2026-08-18T23:16:44Z

+   concatenation and gives each document its own attention-mask ID and
+   zero-based position IDs.
+5. `packed_sequence_thd_vlm_collater` emits `qkv_format: thd`, `seq_lens`, and
+   `seq_lens_padded`; Transformer Engine derives `cu_seqlens` from these values.


Suggested change

`seq_lens_padded`; Transformer Engine derives `cu_seqlens` from these values.

`seq_lens_padded`. Transformer Engine derives `cu_seqlens` from these values.

jgerh · 2026-08-18T23:16:57Z

+
+## Config
+
+Smoke config:


Suggested change

Smoke config:

Use the following smoke config:

jgerh · 2026-08-18T23:17:10Z

+examples/vlm_finetune/stepfun/step3p7_single_node_8xh200_text_sft_offload_8k_thd_packed_bs32_smoke.yaml
+```
+
+Full-data config:


Suggested change

Full-data config:

Use the following full-data config:

jgerh · 2026-08-18T23:17:20Z

+examples/vlm_finetune/stepfun/step3p7_single_node_8xh200_text_sft_offload_8k_thd_packed_bs32_train.yaml
+```
+
+Key settings:


Suggested change

Key settings:

The key settings are:

jgerh · 2026-08-18T23:17:41Z

+
+`local_batch_size` counts packs, not original conversations.
+
+## Validation evidence


Suggested change

## Validation evidence

## Validation Evidence

jgerh · 2026-08-18T23:17:58Z

+had shape `(4, 8192)`, and a multi-document pack emitted sequence lengths
+`[7377, 312, 191, 136]`.
+
+GPU forward/backward/optimizer validation must be performed only after the


Suggested change

GPU forward/backward/optimizer validation must be performed only after the

Perform GPU forward, backward, and optimizer validation only after the

HuiyingLi · 2026-08-18T23:57:44Z

Thank for the PR @new-TonyWang . Looks like there are a lot of things happening in this PR. Would you mind splitting it into several manageable pieces? e.g. I see thd+pp support with step37, cpu offload, dion related, dataset related etc. Thank you so much appreciate it!

new-TonyWang · 2026-08-19T04:20:51Z

@jgerh @HuiyingLi Thank you all, bro! This PR isn’t finished yet. I’ll resolve the merge conflicts, run more compatibility tests, and improve the documentation.

fix(vlm): support packed THD Step3.7 pipeline training

303194a

new-TonyWang requested review from a team and jgerh as code owners August 18, 2026 08:09

nemo-automation-bot Bot added the community-request label Aug 18, 2026

svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Aug 18, 2026

jgerh reviewed Aug 18, 2026

View reviewed changes

		@@ -0,0 +1,63 @@
		# Step-3.7 text SFT with THD packed sequences

		@@ -0,0 +1,63 @@
		# Step-3.7 text SFT with THD packed sequences

		AutoModel's VLM packing pipeline can now consume tokenizer-aware datasets such

	AutoModel's VLM packing pipeline can now consume tokenizer-aware datasets such
	NeMo AutoModel's VLM packing pipeline can now consume tokenizer-aware datasets such

	`seq_lens_padded`; Transformer Engine derives `cu_seqlens` from these values.
	`seq_lens_padded`. Transformer Engine derives `cu_seqlens` from these values.


		`local_batch_size` counts packs, not original conversations.

		## Validation evidence

	GPU forward/backward/optimizer validation must be performed only after the
	Perform GPU forward, backward, and optimizer validation only after the

Uh oh!

Conversation

new-TonyWang commented Aug 18, 2026

Summary

Validation

Notes

Uh oh!

copy-pr-bot Bot commented Aug 18, 2026

Uh oh!

HuiyingLi commented Aug 18, 2026

Uh oh!

jgerh left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HuiyingLi commented Aug 18, 2026

Uh oh!

new-TonyWang commented Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants