[yzhautouskay]

What does this PR do?

Preserve the loader’s chosen dtype for replicated TP parameters by passing dtype=param.dtype to Accelerate. Previously, Accelerate recast weights intended to stay in FP32 to the destination meta parameter’s FP16/BF16 dtype.

Adds CPU regression tests using real safetensors loading and Accelerate. dtype0 is FP16; dtype1 is BF16.

Before:

test_tensor_parallel_loading_preserves_dtype[None-dtype0] PASSED [ 25%]
test_tensor_parallel_loading_preserves_dtype[None-dtype1] PASSED [ 50%]
test_tensor_parallel_loading_preserves_dtype[keep_in_fp32_modules1-dtype0] FAILED [ 75%]
test_tensor_parallel_loading_preserves_dtype[keep_in_fp32_modules1-dtype1] FAILED [100%]

# FP16 loading with an FP32 exception:
>   assert param.dtype == expected_dtype
E   assert torch.float16 == torch.float32

# BF16 loading with an FP32 exception:
>   assert param.dtype == expected_dtype
E   assert torch.bfloat16 == torch.float32

=================== 2 failed, 2 passed, 5 warnings in 2.08s ====================

After:

test_tensor_parallel_loading_preserves_dtype[None-dtype0] PASSED [ 25%]
test_tensor_parallel_loading_preserves_dtype[None-dtype1] PASSED [ 50%]
test_tensor_parallel_loading_preserves_dtype[keep_in_fp32_modules1-dtype0] PASSED [ 75%]
test_tensor_parallel_loading_preserves_dtype[keep_in_fp32_modules1-dtype1] PASSED [100%]

======================== 4 passed, 5 warnings in 0.95s =========================

make style fix-copies passed.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[yzhautouskay]

@sayakpaul Could you please review, based on #15137 (comment)

[github-actions]

Hi @yzhautouskay, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[sayakpaul]

Can this go to tests/models/testing_utils/parallelism.py in a separate class please?

[sayakpaul]

I don't see any problem with the PR. It mimics what we do in the non TP path. So it should really be safe. Cc: @JingyaHuang

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ZihaoW123]

Could the regression include a replicated integer buffer alongside the floating parameters? The loader intentionally casts only floating tensors, and this explicit dtype forwarding should preserve integer checkpoint buffers even when the requested model dtype is FP16/BF16.

[JingyaHuang]

@sayakpaul yes, lgtm.