ms-swift 模型构造与运行时故障分析

模型目录 / 环境变量 / CLI
        │
        ▼
Config resolution
        │  决定模型分支与模块拓扑
        ▼
Module construction
        │  创建 Parameter、persistent buffer、non-persistent buffer
        ▼
post_init() / initialize_weights()
        │  建立构造完成后的确定状态
        ▼
Checkpoint loading
        │  只恢复 checkpoint 中实际序列化的内容
        ▼
Template / Dataset / Collator
        │  生成 input_ids、position_ids、attention_mask、cu_seqlens
        ▼
Forward / Attention / GDN kernel
        │  消费 tensor 与 runtime metadata
        ▼
Generation
        │  EOS、PAD、stopping criteria 决定终止
        ▼
Output

1
2
3
4
5
6
7
8def get_mcore_model(args, hf_config):
    bridge_backend = args.bridge_backend
    if bridge_backend == "megatron-bridge":
        return _get_megatron_bridge_model(args, hf_config)

    config = get_mcore_model_config(args, hf_config)
    models = _get_mcore_model(config)
    return models

ms-swift get_mcore_model()
        ↓
mcore-bridge 解析自定义 ModelMeta
        ↓
构造 Megatron LLM 与 Hugging Face multimodal submodule
        ↓
Qwen3OmniMoeAudioEncoder._from_config(audio_config)
        ↓
SinusoidsPositionEmbedding.__init__()
        ↓
在 meta device 上按公式构造 positional_embedding 的 shape/dtype
        ↓
register_buffer(..., persistent=False)
        ↓
to_empty() / empty materialization 分配未初始化 storage
        ↓
桥接 checkpoint 中的持久参数；该 buffer 不在其中
        ↓
AudioEncoder forward 消费错误的位置表

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28class SinusoidsPositionEmbedding(nn.Module):
    def __init__(self, length, channels, max_timescale=10000):
        super().__init__()
        self.length = length
        self.channels = channels
        self.max_timescale = max_timescale
        position_embedding = (
            self.compute_default_singular_positional_embedding()
        )
        self.register_buffer(
            "positional_embedding",
            position_embedding,
            persistent=False,
        )

    def compute_default_singular_positional_embedding(self):
        log_increment = np.log(self.max_timescale) / (self.channels // 2 - 1)
        inv_timescales = torch.exp(
            -log_increment * torch.arange(self.channels // 2).float()
        )
        scaled_time = (
            torch.arange(self.length)[:, np.newaxis]
            * inv_timescales[np.newaxis, :]
        )
        return torch.cat(
            [torch.sin(scaled_time), torch.cos(scaled_time)],
            dim=1,
        )

1
2
3
4
5
6
7
8
9
10
11positional_embedding = (
    self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]
    .unsqueeze(0)
    .to(padded_embed.dtype)
)
padded_embed = padded_embed + positional_embedding
hidden_states = torch.index_select(
    padded_embed.reshape(-1, padded_embed.shape[-1]),
    0,
    valid_indices,
)

meta 构造只保留 shape 与 dtype
        ↓
empty materialization 分配未初始化 storage
        ↓
checkpoint 不含该 key
        ↓
load_state_dict() 不加载，也不报告 missing key
        ↓
首次 AudioEncoder forward 读取错误位置表

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15@torch.no_grad()
def _init_weights(self, module):
    super()._init_weights(module)

    if isinstance(module, Qwen3OmniMoeCode2Wav):
        init.copy_(
            module.code_offset,
            torch.arange(module.config.num_quantizers).view(1, -1, 1)
            * module.config.codebook_size,
        )
    elif isinstance(module, SinusoidsPositionEmbedding):
        position_embeddings = (
            module.compute_default_singular_positional_embedding()
        )
        init.copy_(module.positional_embedding, position_embeddings)

__init__() 生成确定值
        ↓
post_init() 统一执行模型初始化协议
        ↓
model-specific _init_weights() 重建 non-persistent state
        ↓
checkpoint 加载持久参数
        ↓
普通 from_pretrained 与 bridge _from_config 路径得到相同最终状态

padding_free = true

input_ids.shape = [1, 58000]
cu_seqlens = [0, 2000, 10000, 26000, 58000]
实际参与计算：58000 token


padding_free = false

input_ids.shape = [4, 32000]
四条样本全部补齐到 32000
实际参与计算：128000 token

关闭 padding_free
        ↓
计算规模从 sum(actual_seq_len)
变为 batch_size × max_seq_len
        ↓
padding token 进入每一层 GDN 与 FFN
        ↓
activation + workspace + backward saved tensor 同时增长
        ↓
后续某个普通 buffer 分配成为最终 OOM 触发点

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22chunk_kwargs = {
    "g": g,
    "beta": beta,
    "initial_state": None,
    "output_final_state": cache_params is not None,
    "use_qk_l2norm_in_kernel": True,
}

if (
    cu_seqlens is not None
    and selected_chunk_gated_delta_rule is chunk_gated_delta_rule
):
    chunk_kwargs["cu_seqlens"] = cu_seqlens

core_attn_out, last_recurrent_state = (
    selected_chunk_gated_delta_rule(
        query,
        key,
        value,
        **chunk_kwargs,
    )
)

1
2
3
4
5
6
7
8if (
    _has_multiple_sequences(cu_seqlens)
    and (causal_conv1d is None or chunk_gated_delta_rule is None)
):
    raise ImportError(
        "Qwen3.5 linear attention packing/padding-free with "
        "multiple sequences requires flash-linear-attention"
    )

1
2
3
4
5
6
7
8
9
10
11def _check_padding_free(args, config):
    if not args.padding_free:
        return

    attention_backend = config.attention_backend.name
    if attention_backend == "unfused":
        logger.warning(
            "Attention backend unfused is not supported in "
            "padding-free mode. Setting args.padding_free to False."
        )
        args.padding_free = False

前面各层已经保留大量 activation
        ↓
当前 GDN 已创建 w、u、A 与 chunk state h
        ↓
v_new 再申请一份与 u 同尺寸的 tensor
        ↓
剩余显存不足，错误在这一行暴露

1
2
3
4
5
6
7
8hidden_states, _ = self.self_attn(
    hidden_states=hidden_states,
    attention_mask=attention_mask,
    position_ids=position_ids,
    past_key_values=past_key_values,
    position_embeddings=position_embeddings,
    **kwargs,
)

1
2
3
4
5
6
7
8
9
10
11
12
13
14def _is_packed_sequence(position_ids, batch_size):
    if position_ids is None:
        return False

    increasing_position_sequences = (
        torch.arange(
            position_ids.shape[1],
            device=position_ids.device,
        )
        + position_ids.min()
    )
    return batch_size == 1 and (
        increasing_position_sequences - position_ids
    ).abs().sum().bool()

position_ids.shape = [3, 8, S]
query batch_size = 8
        ↓
batch_size == 1 为 False
        ↓
不进入 position_ids → packed metadata 路径

position_ids.shape = [3, 1, S]
query batch_size = 1
        ↓
_is_packed_sequence() 返回 True
        ↓
三维 M-RoPE position_ids 被误作二维 packed position metadata

1
2
3
4
5
6
7
8
9
10
11
12
13
14def prepare_fa_kwargs_from_position_ids(position_ids):
    position_ids = position_ids.reshape(-1)
    indices_q = (position_ids == 0).nonzero().view(-1)
    cu_seq_lens_q = torch.cat(
        (
            indices_q.to(dtype=torch.int32),
            torch.tensor(
                position_ids.size(),
                dtype=torch.int32,
                device=position_ids.device,
            ),
        )
    )
    return cu_seq_lens_q

真实 query
q.shape = [256, num_heads, head_dim]
q token count = 256

错误输入 metadata
position_ids.shape = [3, 1, 256]
position_ids.reshape(-1).shape = [768]
cu_seqlens_q = [0, 256, 512, 768]

关键不变量被破坏
cu_seqlens_q[-1] = 768
q.shape[0] = 256

尾批 B=1
        ↓
FA2 packed-sequence 探测被激活
        ↓
3D M-RoPE position_ids 被误作 packed metadata
        ↓
cu_seqlens[-1] != q token count
        ↓
FA2 varlen kernel 越界或产生错误 attention
        ↓
logits / EOS 行为退化
        ↓
CUDA 异步错误在后续同步点暴露

1
2
3
4
5
6
7
8if cu_seqlens_q is not None:
    # q 是即将送入 varlen kernel 的 flattened query operand
    if int(cu_seqlens_q[-1]) != q.shape[0]:
        raise ValueError(
            "FlashAttention metadata mismatch: "
            f"cu_seqlens_q[-1]={int(cu_seqlens_q[-1])}, "
            f"query_tokens={q.shape[0]}"
        )

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15# Transformers 5.2.0 / 5.3.0
hidden_states = decoder_layer(
    hidden_states,
    position_embeddings=position_embeddings,
    position_ids=position_ids,       # [3, B, S]
    ...,
)

# Transformers >= 5.4.0
hidden_states = decoder_layer(
    hidden_states,
    position_embeddings=position_embeddings,
    position_ids=text_position_ids,  # [B, S]
    ...,
)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23class Qwen3OmniMoeForConditionalGeneration(
    Qwen3OmniMoePreTrainedModel,
    GenerationMixin,
):
    def __init__(self, config):
        super().__init__(config)
        self.thinker = (
            Qwen3OmniMoeThinkerForConditionalGeneration
            ._from_config(config.thinker_config)
        )
        self.has_talker = config.enable_audio_output
        if self.has_talker:
            self.enable_talker()
        self.post_init()

    def enable_talker(self):
        self.talker = (
            Qwen3OmniMoeTalkerForConditionalGeneration
            ._from_config(self.config.talker_config)
        )
        self.code2wav = Qwen3OmniMoeCode2Wav._from_config(
            self.config.code2wav_config
        )

读取 config.json
        ↓
创建 Thinker
        ↓
检查 enable_audio_output
        ├── false → 模型仅包含 Thinker
        │
        └── true  → enable_talker()
                        ├── 创建 Talker
                        └── 创建 Code2Wav

1
2
3
4
5
6
7
8
9
10class Qwen3OmniLoader(ModelLoader):
    def get_config(self, model_dir: str):
        self._check_qwen_omni_utils()
        config = super().get_config(model_dir)
        enable_audio_output = get_env_args(
            "ENABLE_AUDIO_OUTPUT", bool, None
        )
        if enable_audio_output is not None:
            config.enable_audio_output = enable_audio_output
        return config

1ENABLE_AUDIO_OUTPUT=false swift megatron sft ...

1
2
3
4model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(
    model_path,
    enable_audio_output=False,
)