Skip to content

Integrate Lineage DeepSeek-V3 model (dsv3.py) into MaxText. #5072

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Integrate Lineage DeepSeek-V3 model (dsv3.py) into MaxText. #5072

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/maxtext/configs/base.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,11 @@ topk_routing_group: -1 # number of top groups to route inputs. For EP,
# all-to-all communication with compute. Currently only implemented with DeepSeek sparse layers.
use_batch_split_schedule: false # a flag if splitting batch into micro-batches to hide communications that yields performance benefits.
batch_split_factor: 1 # the factor by which to split the batch. Only used if use_batch_split_schedule is true.
use_lineage: false # a flag to use Lineage DeepSeek-V3 execution.
lineage_attention_sharding: "head" # attention sharding strategy for Lineage ('head' or 'sequence').
lineage_activation_checkpointing: true # a flag to use activation checkpointing in Lineage layers.
lineage_capacity_factor: 0.5 # positive capacity factor determining the destination buffer size for Lineage sparse dispatch. If <= 0, falls back to capacity_factor or ragged_buffer_factor.
lineage_mesh_axes_mapping: {} # custom mapping from Lineage logical axes to mesh physical axes.

# For complex architectures like llama4 there are repeated sets of
# inhomogeneous layers. E.g. maverick uses [dense+rope, moe+rope, dense+rope, moe+nope]
Expand Down
37 changes: 37 additions & 0 deletions src/maxtext/configs/types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,6 +236,7 @@ class ProfilerType(str, Enum):
"deepseek3-671b",
"deepseek3-671b-2dfsdp",
"deepseek3-671b-batchsplit",
"deepseek3-671b-lineage",
"deepseek3-test",
"deepseek3-tiny",
"deepseek3.2-671b",
Expand Down Expand Up @@ -1127,6 +1128,42 @@ class DeepSeekMoE(BaseModel):
1,
description="Factor by which to split the batch into micro-batches. Only used if use_batch_split_schedule is True.",
)
use_lineage: bool = Field(
False,
description="Whether to use Lineage DeepSeek-V3 execution.",
)
lineage_attention_sharding: Literal["head", "sequence"] = Field(
"head",
description=("Attention sharding strategy for Lineage ('head' or 'sequence')."),
)
lineage_activation_checkpointing: bool = Field(
True,
description="Whether to use activation checkpointing in Lineage layers.",
)
lineage_capacity_factor: float = Field(
0.5,
description=(
"Positive capacity factor determining the destination buffer size for"
" Lineage sparse dispatch. If <= 0, falls back to capacity_factor or"
" ragged_buffer_factor."
),
)
lineage_mesh_axes_mapping: dict[str, Any] = Field(
default_factory=dict,
description=("Custom mapping from Lineage logical axes to mesh physical axes."),
)

@model_validator(mode="after")
def validate_lineage(self) -> "DeepSeekMoE":
"""Validates that Lineage DeepSeek-V3 execution requirements are met."""
if self.use_lineage:
scan_layers = getattr(self, "scan_layers", None)
if scan_layers is not None and not scan_layers:
raise ValueError("use_lineage=True requires scan_layers=True.")
decoder_block = getattr(self, "decoder_block", None)
if decoder_block is not None and decoder_block != DecoderBlockType.DEEPSEEK:
raise ValueError(f"use_lineage=True requires decoder_block='deepseek', got decoder_block={decoder_block!r}.")
return self


class Qwen3Next(BaseModel):
Expand Down
249 changes: 150 additions & 99 deletions src/maxtext/layers/decoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,11 @@
qwen3_5,
simple_layer,
)

try:
from maxtext.models import lineage_adapter
except ImportError:
lineage_adapter = None
from maxtext.multimodal import utils as mm_utils
from maxtext.utils.sharding import create_sharding
from maxtext.utils import max_logging
Expand Down Expand Up @@ -958,105 +963,19 @@ def __call__(
else:
if cfg.scan_layers:
if cfg.decoder_block == DecoderBlockType.DEEPSEEK:
assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek."
layer_call_kwargs = {
"previous_chunk": previous_chunk,
"slot": slot,
}
dense_layer = RemattedBlockLayers[0]
moe_layer = RemattedBlockLayers[1]
if cfg.engram_layers:
original_dense_call = dense_layer.__call__
original_moe_call = moe_layer.__call__
dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs)
moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs)

common_kwargs = {
"dense_layer": dense_layer,
"moe_layer": moe_layer,
"original_dense_call": original_dense_call,
"original_moe_call": original_moe_call,
"layer_call_kwargs": layer_call_kwargs,
"decoder_segment_ids": decoder_segment_ids,
"decoder_positions": decoder_positions,
"deterministic": deterministic,
"model_mode": model_mode,
"decoder_input_tokens": decoder_input_tokens,
"broadcast_args": broadcast_args,
}

# Apply Dense Layers
y = self._apply_interleaved_scanned_layers(
y,
layer_type="dense",
start_idx=0,
end_idx=cfg.first_num_dense_layers,
engram_indices=cfg.engram_layers,
**common_kwargs,
)

# Apply MoE Layers
y = self._apply_interleaved_scanned_layers(
y,
layer_type="moe",
start_idx=cfg.first_num_dense_layers,
end_idx=cfg.num_decoder_layers,
engram_indices=cfg.engram_layers,
**common_kwargs,
)
else:
dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs)
y, _ = self.scan_decoder_layers(
cfg,
dense_layer,
cfg.first_num_dense_layers,
"dense_layers",
mesh,
in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
model_mode=model_mode,
)(y, *broadcast_args)
moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs)
num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers

# If batch-split schedule is used and initialization is complete,
# as detected by immutable params, use deepseek_batchsplit custom
# scan with initialized parameters.
if cfg.use_batch_split_schedule and not self.is_mutable_collection("params"):
# old version of batch-split that fully uses qwix quantization.
if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization:
y = deepseek_batchsplit_fp8.scan_batch_split_layers(
y,
self.variables["params"]["moe_layers"],
decoder_positions,
decoder_segment_ids,
model_mode=model_mode,
mesh=mesh,
quant=self.quant,
cfg=cfg,
policy=policy,
)
else:
# bf16 and fp8 code path for pure-JAX batch-split.
# fp8 code path supports both manual quantization and qwix
# quantization.
y = deepseek_batchsplit.scan_batch_split_layers(
y,
self.variables["params"]["moe_layers"],
decoder_positions,
mesh=mesh,
cfg=cfg,
num_layers=num_moe_layers,
)
else:
y, _ = self.scan_decoder_layers(
cfg,
moe_layer,
num_moe_layers,
"moe_layers",
mesh,
in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
model_mode=model_mode,
)(y, *broadcast_args)
y = self._apply_deepseek_scanned_blocks(
y,
RemattedBlockLayers,
decoder_segment_ids,
decoder_positions,
deterministic,
model_mode,
decoder_input_tokens,
previous_chunk,
slot,
broadcast_args,
policy,
)
elif cfg.decoder_block == DecoderBlockType.GEMMA3:
bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None
y = self._apply_gemma3_scanned_blocks(
Expand Down Expand Up @@ -1344,6 +1263,138 @@ def __call__(
# and the raw hidden state needed for auxiliary tasks.
return logits, hidden_state, kv_caches

def _apply_deepseek_scanned_blocks(
self,
y,
RemattedBlockLayers,
decoder_segment_ids,
decoder_positions,
deterministic,
model_mode,
decoder_input_tokens,
previous_chunk,
slot,
broadcast_args,
policy,
):
"""Applies DeepSeek scanned decoder blocks, handling dense and MoE layers."""
cfg = self.config
mesh = self.mesh

assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek."
layer_call_kwargs = {
"previous_chunk": previous_chunk,
"slot": slot,
}
dense_layer = RemattedBlockLayers[0]
moe_layer = RemattedBlockLayers[1]
if cfg.engram_layers:
original_dense_call = dense_layer.__call__
original_moe_call = moe_layer.__call__
dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs)
moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs)

common_kwargs = {
"dense_layer": dense_layer,
"moe_layer": moe_layer,
"original_dense_call": original_dense_call,
"original_moe_call": original_moe_call,
"layer_call_kwargs": layer_call_kwargs,
"decoder_segment_ids": decoder_segment_ids,
"decoder_positions": decoder_positions,
"deterministic": deterministic,
"model_mode": model_mode,
"decoder_input_tokens": decoder_input_tokens,
"broadcast_args": broadcast_args,
}

# Apply Dense Layers
y = self._apply_interleaved_scanned_layers(
y,
layer_type="dense",
start_idx=0,
end_idx=cfg.first_num_dense_layers,
engram_indices=cfg.engram_layers,
**common_kwargs,
)

# Apply MoE Layers
y = self._apply_interleaved_scanned_layers(
y,
layer_type="moe",
start_idx=cfg.first_num_dense_layers,
end_idx=cfg.num_decoder_layers,
engram_indices=cfg.engram_layers,
**common_kwargs,
)
else:
num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers
if getattr(cfg, "use_lineage", False) and not self.is_mutable_collection("params"):
y = lineage_adapter.run_lineage_dsv3(
inputs=y,
dense_params=self.variables["params"]["dense_layers"],
sparse_params=self.variables["params"]["moe_layers"],
decoder_positions=decoder_positions,
mesh=mesh,
cfg=cfg,
decoder_segment_ids=decoder_segment_ids,
num_dense_layers=cfg.first_num_dense_layers,
num_sparse_layers=num_moe_layers,
)
else:
dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs)
y, _ = self.scan_decoder_layers(
cfg,
dense_layer,
cfg.first_num_dense_layers,
"dense_layers",
mesh,
in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
model_mode=model_mode,
)(y, *broadcast_args)
moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs)

# If batch-split schedule is used and initialization is complete,
# as detected by immutable params, use deepseek_batchsplit custom
# scan with initialized parameters.
if cfg.use_batch_split_schedule and not self.is_mutable_collection("params"):
# old version of batch-split that fully uses qwix quantization.
if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization:
y = deepseek_batchsplit_fp8.scan_batch_split_layers(
y,
self.variables["params"]["moe_layers"],
decoder_positions,
decoder_segment_ids,
model_mode=model_mode,
mesh=mesh,
quant=self.quant,
cfg=cfg,
policy=policy,
)
else:
# bf16 and fp8 code path for pure-JAX batch-split.
# fp8 code path supports both manual quantization and qwix
# quantization.
y = deepseek_batchsplit.scan_batch_split_layers(
y,
self.variables["params"]["moe_layers"],
decoder_positions,
mesh=mesh,
cfg=cfg,
num_layers=num_moe_layers,
)
else:
y, _ = self.scan_decoder_layers(
cfg,
moe_layer,
num_moe_layers,
"moe_layers",
mesh,
in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
model_mode=model_mode,
)(y, *broadcast_args)
return y

def _apply_gemma3_scanned_blocks(
self,
y,
Expand Down
Loading
Loading