fix: load checkpoints in weights-only mode

[jmccaffrey-nv]

Summary

• make weights-only behavior explicit for local, Hugging Face cache, S3, and distributed-checkpoint cache reads
• preserve safetensors handling and tensor state-dict compatibility for .pt, .pth, and .ckpt files
• cover non-weight objects and normal tensor-only checkpoints across each supported path

Rationale

The checkpoint-loading API returns state dictionaries. Loading legacy formats in weights-only mode keeps that contract consistent across supported PyTorch versions without removing existing checkpoint formats.

Validation

• focused checkpoint suite: 9 passed
• changed-file lint and formatting checks passed
• changed-file type checks passed
• broader CPU suite: 677 passed, with two unrelated environment-only cases excluded

[copy-pr-bot]

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.

[greptile-apps]

Confidence Score: 4/5

The PR appears safe to merge, with a non-blocking test-coverage gap for valid S3 and distributed-cache checkpoint reads.

Findings

1.  Missing success-path coverage ▶

Summary

• Changes three torch.load call sites to use weights_only=True.
• Adds coverage proving non-weight objects are rejected without executing their reducers.
• Verifies local .pt, .pth, and .ckpt tensor state dictionaries still load.
• Leaves successful tensor-only S3 and distributed-cache reads without direct regression coverage.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    Local[Local checkpoint] --> SafeLoad[torch.load weights_only=True]
    HF[Hugging Face cached file] --> SafeLoad
    S3[S3 checkpoint bytes] --> SafeLoad
    DCP[Distributed checkpoint cache] --> SafeLoad
    SafeLoad --> StateDict[Tensor state dictionary]
    StateDict --> Model[Model load_state_dict]
    SafeLoad -. rejects .-> Object[Executable non-weight object]

Loading

[jmccaffrey-nv]

/ok to test aae515d