[model] Support bailing_hybrid

[Jintao-Huang]

https://huggingface.co/inclusionAI/Ling-2.6-flash/discussions/6/files

[gemini-code-assist]

Code Review

This pull request introduces support for the bailing_hybrid model, including its configuration mapping and a specialized loader that handles hybrid attention layers. Review feedback highlights the need for safer and more efficient logic when retrieving transformer layer specifications, specifically recommending a try...finally block to ensure configuration state is restored. Additionally, it was suggested to remove redundant method overrides in the LinearAttention class and clean up several unused imports.

[gemini-code-assist]

The current implementation for getting linear_layer_specs by temporarily modifying self.config.multi_latent_attention has a couple of issues:

1. Safety: If super().get_transformer_layer_spec() raises an exception, self.config.multi_latent_attention will not be restored to its original value. This could lead to unexpected behavior in subsequent operations. Using a try...finally block is recommended for safety.
2. Efficiency: super().get_transformer_layer_spec() is called twice. If self.config.multi_latent_attention is False to begin with, both calls are identical, which is redundant and inefficient.

Consider refactoring this logic to be safer and more efficient.

Suggested change

	layer_specs = super().get_transformer_layer_spec(vp_stage=vp_stage)
	multi_latent_attention = self.config.multi_latent_attention
	self.config.multi_latent_attention = False
	linear_layer_specs = super().get_transformer_layer_spec(vp_stage=vp_stage)
	self.config.multi_latent_attention = multi_latent_attention
	multi_latent_attention = self.config.multi_latent_attention
	if multi_latent_attention:
	layer_specs = super().get_transformer_layer_spec(vp_stage=vp_stage)
	try:
	self.config.multi_latent_attention = False
	linear_layer_specs = super().get_transformer_layer_spec(vp_stage=vp_stage)
	finally:
	self.config.multi_latent_attention = multi_latent_attention
	else:
	layer_specs = super().get_transformer_layer_spec(vp_stage=vp_stage)
	linear_layer_specs = layer_specs

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for the bailing_hybrid model, including its configuration, model definition, and custom linear attention layers. It also updates the documentation and refactors the rotary position embedding patching logic to be more flexible. A critical issue was identified in the newly added bailing_hybrid.py file, where comparing PyTorch tensors directly using == in an assertion can lead to a runtime error. A code suggestion has been provided to resolve this issue safely using torch.equal.

[Jintao-Huang]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for the bailing_hybrid model, implementing LinearAttention using flash-linear-attention, refactoring bailing_moe to simplify its QKV weight layout handling, and updating configuration parsing and documentation. The review feedback highlights several critical issues: potential runtime crashes in bailing_hybrid.py from using .view() on non-contiguous tensors, unpacking a None rotary embedding, division by zero when num_layers is 1, and using an in-place sigmoid on a tensor requiring gradients. Additionally, the reviewer noted a runtime crash in bailing_moe.py when performing dist.all_reduce on a boolean tensor, and a potential bug in patcher.py where cp_group could be missed if passed positionally.

[gemini-code-assist]

If rotary_pos_emb is None (e.g., when rotary embedding is disabled or not used), unpacking it directly will raise a TypeError: cannot unpack non-iterable NoneType object. Adding a guard prevents this crash.

Suggested change

	q_pos_emb, k_pos_emb = rotary_pos_emb
	q_pos_emb, k_pos_emb = rotary_pos_emb if rotary_pos_emb is not None else (None, None)