`_is_hf_initialized` per-param check doesn't skip redundant init for built-in models on non-CUDA hardware

[gygdh-001]

System Info

• transformers version: 4.57.6
• Platform: Ascend NPU 910B, CANN 8.5.0
• Python version: 3.11
• PyTorch version: 2.9.0, torch_npu 2.9.0
• Hardware: 16 × Ascend 910B, FSDP2

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

# Minimal reproduction: any built-in model under FSDP on Ascend NPU
# The _is_hf_initialized per-param flags are completely ignored for non-custom-code models

from transformers import AutoModel
import torch
import torch.distributed as dist
import torch_npu

# 1. Init FSDP process group
dist.init_process_group(backend='hccl')

# 2. Load a built-in model (NOT trust_remote_code)
model = AutoModel.from_pretrained("google/umt5-small")

# 3. Under FSDP2, _initialize_missing_keys() sets _is_hf_initialized = True
#    on all params of non-rank0 processes.
#    But _initialize_weights() NEVER reads those flags because is_custom_code=False.

# 4. Result: all non-rank0 ranks call _init_weights() → normal_() on every param.
#    On Ascend NPU, normal_ is ~1000x slower than CUDA (~1s vs ~1ms per large tensor).
#    For a large model: 219 redundant normal_ calls × 16 ranks = ~224s wasted.

# Root cause location: modeling_utils.py L2564-2574
# _check_per_param_flag = is_custom_code  ← built-in models always False here

Expected behavior

Root Cause

_initialize_weights() in modeling_utils.py has a per-parameter _is_hf_initialized check at L2564-2574. The variable _check_per_param_flag is set to is_custom_code, which is False for all built-in models (UMT5, BERT, T5, etc.). This makes the per-param flag check dead code for those models — even though _initialize_missing_keys() at L4882 correctly sets the flags on all params of non-rank0 FSDP processes.

Why it matters on Ascend NPU

torch.nn.init.normal_ on Ascend NPU is ~1000x slower than CUDA (~1s vs ~1ms per large tensor). Each call triggers a full kernel compilation and launch. Under FSDP2, every non-rank0 process redundantly initializes all params via _init_weights() → normal_(), even though those params will be overwritten by rank0's broadcast.

Measured Impact (MOVA training, 16 × Ascend 910B)

Metric	Before fix	After fix	Delta
Training launch	463s	136s	−71%
Redundant normal_ calls skipped	—	219 per rank	—

Precision Validation

Verified that skipping redundant _init_weights on already-initialized params has zero precision impact:

• 16 ranks × 4397 parameters: weight fingerprints all bit-level identical
• 2-step loss: MAE=0.000e+00 across all 16 ranks for loss / v_loss / a_loss

Proposed Fix

Minimal, NPU-scoped (PR #47335 v3):

# modeling_utils.py L2566
_check_per_param_flag = is_custom_code
if not _check_per_param_flag:
    try:
        import torch_npu
        _check_per_param_flag = True
    except ImportError:
        pass

is_torch_npu_available() → False when torch_npu not installed — zero impact on CUDA/ROCm/XPU.

Or more generally, just remove the is_custom_code gate. The per-param flag is only set by _initialize_missing_keys (FSDP path), so skipping _init_weights for flagged params should be safe for all model types.

Related

• PR Remove is_custom_code gate from per-param _is_hf_initialized check in _initialize_weights #47335: fix with is_torch_npu_available() gate (v3), reviewed by @Rocketknight1
• Ascend NPU maintainer: @ivarflakstad

cc @Rocketknight1 @ivarflakstad

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.