[yingguo-trt]

What does this PR do?

Type of change: Bug fix.

LTX-2 distillation with training_mode: full names its standalone weights .safetensors, but Accelerator.save() defaults to PyTorch serialization. Loading that file with safetensors fails with Error while deserializing header: header too large.

Pass safe_serialization=True so the saved contents match the advertised format, and document the fix in the changelog.

Usage

No configuration changes are required. Full-mode training now writes valid safetensors weights. Existing incorrectly serialized files are not converted.

Testing

• Reused the existing external QA case test_ltx_full_checkpoint_is_safetensors; no new tests added. It executes the production save method extracted unchanged from the example, with a tiny CPU model and real Accelerate/safetensors APIs. The baseline fails with the header error; the fix passes with exact key and BF16 tensor equality. This is method-level validation, not a full LTX training run.
• Environment: Python 3.12.10, PyTorch 2.12.0+cpu, Accelerate 1.15.0, safetensors 0.8.0, and ModelOpt 0.48.0rc1 for library imports; example source tested before and after this change.
• Applicable pre-commit hooks and git diff --check pass. The mypy hook environment passes with --platform linux; the pre-existing Windows os.sched_getaffinity error also reproduces on the unmodified source and was excluded.
• Full LTX GPU/FSDP training and inference were not run.

Before your PR is "Ready for review"

• Is this change backward compatible?: Yes for the documented safetensors output contract; no API or configuration changes.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A.
• Did you write any new necessary tests?: N/A; reused the existing external QA regression described above. No in-repository regression test added.
• Did you update Changelog?: Yes.
• Did you get Claude approval on this PR?: Pending.

Additional Information

NVBug 6982528.

Summary by CodeRabbit

• Bug Fixes
  • Fixed the format used when saving full-mode LTX-2 distillation checkpoints. Checkpoint files now use the format indicated by their file extension, improving compatibility with tools that load them and ensuring consistent handling by compatible loaders. This correction applies to full-mode model checkpoints; LoRA checkpoint saving remains unchanged.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: main
• File: .coderabbit.yaml
• Commit: c8c4389
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 0fa1f2af-b3e8-4c51-b900-baa5af5f0110

Commits

Reviewing files that changed from the base of the PR and between 2f20f21 and 80eca54.

Files selected for processing: 1

• CHANGELOG.rst

Files skipped from review as they are similar to previous changes: 1

• CHANGELOG.rst

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

Non-LoRA distillation checkpoints now use safe serialization when saved. The LoRA save path remains unchanged. The changelog records the LTX-2 full-mode checkpoint format fix.

Changes: Distillation checkpoint format

Layer / File(s)	Summary
Save checkpoints with safe serialization examples/diffusers/distillation/distillation_trainer.py, CHANGELOG.rst	The non-LoRA checkpoint save call enables safe serialization. The changelog records the LTX-2 full-mode checkpoint format fix.

Priority: Normal

Estimated code review effort: 1 (Trivial) | ~5 minutes

Change: Bug fix

Suggested reviewers: kevalmorabia97, mxino

Merge Risk: Minimal · up to 80eca

The changelog entry is in the intended release section. No actionable merge-blocking issue remains after normal checks.

Pre-merge checks: 6

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title "[NVBug 6982528] Save LTX full-mode checkpoints as safetensors" matches the main change. The PR passes safe_serialization=True to Accelerator.save() for LTX-2 distillation full-mode checkpoi…
Docstring Coverage		Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files. (1 skipped: 1 …
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The PR changes only CHANGELOG.rst and one call in examples/diffusers/distillation/distillation_trainer.py. The code change adds safe_serialization=True to Accelerator.save(); it adds no torch.lo…

Full details: Title check

Explanation

The title "[NVBug 6982528] Save LTX full-mode checkpoints as safetensors" matches the main change. The PR passes safe_serialization=True to Accelerator.save() for LTX-2 distillation full-mode checkpoints, so the saved files match their .safetensors extension. The bug tracker prefix is extra detail but does not make the title misleading. The title is specific and short, so it passes.

Full details: Docstring Coverage

Explanation

Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files. (1 skipped: 1 unsupported.)

Full details: Security Anti-Patterns

Explanation

PASS. The PR changes only CHANGELOG.rst and one call in examples/diffusers/distillation/distillation_trainer.py. The code change adds safe_serialization=True to Accelerator.save(); it adds no torch.load(..., weights_only=False), pickle-enabled NumPy load, hardcoded trust_remote_code=True, input-driven eval()/exec(), or # nosec. No dependency manifest changes were made.

Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.51%. Comparing base (c8c4389) to head (80eca54).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3015   +/-   ##
=======================================
  Coverage   77.51%   77.51%           
=======================================
  Files         661      661           
  Lines       74937    74937           
=======================================
  Hits        58086    58086           
  Misses      16851    16851           

Flag	Coverage Δ
examples-diffusers	19.90% <ø> (-0.01%)	⬇️
unit	61.13% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.