[Bugfix] [Offloading] Save disk-offloaded buffers, Save converted weights

[kylesayrs]

Purpose

• Parlay two bugfixes in one for the purpose of enabling the saving of models with disk offloaded weights
  • Fix saving models which are disk offloaded with offload_buffers=True
  • Fix saving models which are disk offloaded with weight conversions

Changes

• Replace get_parameter call with get_parameter_or_buffer call
• Expand load_offloaded_parameter to load a state dict of all weights associated with the checkpoint weight
  • checkpoint weight (all/any) -> model weight (one) -> reverted weights (all)
  • By using this loaded state dict to update the original state dict, we avoid redundant loading of offloaded weights
• Loosen requirements for is_offloaded flag
  • This is backwards compatible safe, as full disk offloaded was never supported in previous releases anyways

Testing

Able to save disk-offloaded models with conversion mappings now

from transformers import AutoModelForCausalLM

# Load model with full disk offload
model = AutoModelForCausalLM.from_pretrained(
    "inference-optimization/DSV4-tiny-empty",
    device_map="auto",
    max_memory={},
    offload_folder="offload_folder",
)

# Save the model
model.save_pretrained("tmp_save")

Used these changes to quantize RedHatAI/DeepSeek-V4-Pro-NVFP4-FP8 and RedHatAI/GLM-5.2-NVFP4-FP8

Suggested Reviewers

• @SunMarc @zucchini-nlp @ArthurZucker

[SunMarc]

Thanks, left a couple of comments but I think @Cyrilvallez might have better ideas on how to deal with those as he's the one who coded this !

[SunMarc]

any specific reason to sort ?

[kylesayrs]

Note that load_offloaded_parameter may load multiple weights for a single tensor.
While it is possible to overload CPU memory by loading parameters in a bad order,
in practice split_torch_state_dict_into_shards preserves weight locality

sorting helps reduce the chances that bad load ordering occurs.

An example of bad load ordering would be

"layers.0.experts.0.up_proj" -> loads "layers.0.experts.gate_up_proj"
"layers.1.experts.0.up_proj" -> loads "layers.1.experts.gate_up_proj"
"layers.2.experts.0.up_proj" -> loads "layers.2.experts.gate_up_proj"

In this scenario, 3 separate gate_up_proj weights have been loaded onto cpu, but only 3 shard weights have been consumed by state_dict.pop.

Sorting reduces the chances that split_torch_state_dict_into_shards gives an adversarially bad ordering. It doesn't fix adversarially bad ordering between shards, but there's not much we can do about that.

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 28566046728:2
Result: failure | Jobs: 14 | Tests: 72,772 | Failures: 1 | Duration: 16h 25m

[Cyrilvallez]

Hey @kylesayrs! I took the liberty to open #47018 to fix the issue, I believe it is simpler and more robust in general!
Let me know if something is still unclear