[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.