[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

@@ -579,12 +579,14 @@ class LLaDA2MoeBlock(nn.Module):
self.is_layer_sparse = self._is_layer_sparse(config, layer_id=layer_id)
is_previous_layer_sparse = self._is_layer_sparse(config, layer_id=layer_id - 1)
is_next_layer_sparse = self._is_layer_sparse(config, layer_id=layer_id + 1)
self.layer_scatter_modes = LayerScatterModes.init_new(
layer_id=layer_id,
num_layers=config.num_hidden_layers,
is_layer_sparse=self.is_layer_sparse,
is_previous_layer_sparse=is_previous_layer_sparse,
is_next_layer_sparse=is_next_layer_sparse,
)
self.is_last_layer = self.layer_id == config.num_hidden_layers - 1