Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsMuse Sparknemo_automodel.components.models.muse_spark.magi_cp

nemo_automodel.components.models.muse_spark.magi_cp

View as Markdown

Model-owned packed Magi CP with shared token ownership across Muse masks.

Module Contents

Classes

NameDescription
MuseMagiPlanPer-microbatch keys; all masks retain the same dispatched token order.
_MagiAttentionKeyMagi key fields consumed by the Muse integration.
_MuseMagiBatchPlansStep-owned plans replicated by PP; a chunked index selects each plan.
_MuseMagiModelMuse model state needed to build model-owned Magi CP inputs.

Functions

NameDescription
_build_magi_planBuild masks with one shared token partition for a packed microbatch.
prepare_magi_inputsPrepare model-owned packed CP, including independent pipeline microbatches.

Data

MuseMagiBatch

MuseMagiBatchValue

logger

API

class nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan(
    keys: dict[int, nemo_automodel.components.models.muse_spark.magi_cp._MagiAttentionKey]
)
Dataclass

Per-microbatch keys; all masks retain the same dispatched token order.

keys
dict[int, _MagiAttentionKey]
nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan.attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    window: int,
    scale: float
) -> torch.Tensor

Compute local THD output with Magi’s own KV communication.

Parameters:

query
torch.Tensor

CP-local queries [tokens, query_heads, head_dim].

key
torch.Tensor

CP-local keys [tokens, kv_heads, head_dim].

value
torch.Tensor

CP-local values [tokens, kv_heads, head_dim].

window
int

Sliding-window length, or zero for full causal attention.

scale
float

Attention score scale.

Returns: torch.Tensor

Attention output [tokens, query_heads, head_dim] in the local token order.

class nemo_automodel.components.models.muse_spark.magi_cp._MagiAttentionKey()
Protocol

Magi key fields consumed by the Muse integration.

total_seqlen_q
int
class nemo_automodel.components.models.muse_spark.magi_cp._MuseMagiBatchPlans(
    plans: tuple[nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan, ...]
)
Dataclass

Step-owned plans replicated by PP; a chunked index selects each plan.

plans
tuple[MuseMagiPlan, ...]
nemo_automodel.components.models.muse_spark.magi_cp._MuseMagiBatchPlans.select(
    microbatch_index: torch.Tensor
) -> nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan

Select without a mutable cursor, before entering checkpointed layers.

Parameters:

microbatch_index
torch.Tensor

Int64 tensor [1], split by PP and possibly moved by FSDP.

Returns: MuseMagiPlan

The immutable microbatch plan retained by that forward’s AC graph.

class nemo_automodel.components.models.muse_spark.magi_cp._MuseMagiModel()
Protocol

Muse model state needed to build model-owned Magi CP inputs.

config
MuseSpark1xOmniConfig
cp_mesh
DeviceMesh | None
nemo_automodel.components.models.muse_spark.magi_cp._MuseMagiModel._is_pipeline_parallel_stage() -> bool

Return whether this module is a physical pipeline stage.

nemo_automodel.components.models.muse_spark.magi_cp._build_magi_plan(
    cu_seqlens: torch.Tensor,
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    group: torch.distributed.ProcessGroup
) -> nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiPlan

Build masks with one shared token partition for a packed microbatch.

Parameters:

cu_seqlens
torch.Tensor

CPU int32 padded document boundaries [documents + 1].

config
MuseSpark1xOmniConfig

Muse attention configuration, shared across pipeline stages.

group
torch.distributed.ProcessGroup

This pipeline stage’s context-parallel process group.

Returns: MuseMagiPlan

Full-causal and sliding-window keys with identical local token ownership.

Prepare model-owned packed CP, including independent pipeline microbatches.

Parameters:

model
_MuseMagiModel

Full Muse model or a local pipeline stage.

batch
MuseMagiBatch

Global tokens/labels/positions [batch, sequence] and packed document lengths [batch, documents]. Each PP microbatch holds one pack.

num_chunks
int

Number of independent pipeline microbatches.

Returns: dict[str, ContextParallelSharder]

A model-owned CP sharder. In PP, token ids stay global [chunks, sequence];

nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiBatch = dict[str, MuseMagiBatchValue]
 
nemo_automodel.components.models.muse_spark.magi_cp.MuseMagiBatchValue = torch.Tensor | str | int | float | bool | MuseMagiPlan | _MuseMagiBatchPlans | N...
 
nemo_automodel.components.models.muse_spark.magi_cp.logger = logging.getLogger(__name__)