[bugfix] fix TBO crashes when attn_tp_size > 1 (#13730)

Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
This commit is contained in:
Yuhao Yao
2025-12-12 10:18:40 +08:00
committed by GitHub
parent 0aa3dec5c7
commit e9e7f15eb5
20 changed files with 285 additions and 16 deletions

View File

@@ -380,6 +380,8 @@ class LongcatFlashDecoderLayer(nn.Module):
num_layers=config.num_hidden_layers,
is_layer_sparse=False,
is_previous_layer_sparse=False,
# TODO: Check if the following is correct.
is_next_layer_sparse=False,
)
for i in range(2)
]
@@ -398,6 +400,8 @@ class LongcatFlashDecoderLayer(nn.Module):
num_layers=config.num_hidden_layers,
is_layer_sparse=True,
is_previous_layer_sparse=True,
# TODO: Check if the following is correct.
is_next_layer_sparse=True,
)
self.moe_layer_communicator = LayerCommunicator(
layer_scatter_modes=self.moe_layer_scatter_modes,