nemo_automodel.components.models.flint.packing

Flint text-token ownership across TE context and pipeline parallelism.

Module Contents

Functions

Name	Description
`_prepare_inputs_for_cp`	Choose a shared text sharder while preserving PP’s global-id boundary.

API

nemo_automodel.components.models.flint.packing._prepare_inputs_for_cp(
    model: nemo_automodel.components.models.flint.model.FlintForCausalLM,
    batch: dict[str, typing.Any],
    num_chunks: int
) -> dict[str, typing.Any]

Choose a shared text sharder while preserving PP’s global-id boundary.

Parameters:

model

FlintForCausalLM

Native model, possibly one pipeline stage, with optional CP mesh.

batch

dict[str, Any]

Global token tensors [B,S] and optional document lengths [B,N].

num_chunks

int

Number of independent pipeline microbatch token spaces.

Returns: dict[str, Any]

A model-owned sharder and padding metadata. Packed local indices have