nemo_automodel.components.models.flint.model

Native full-sequence Flint causal language model with explicit training scope.

Module Contents

Classes

Name	Description
`FlintForCausalLM`	Eager Flint baseline, with protected trainable sinks and untied FP32 head.
`FlintModel`	Untied input embedding, scheduled decoder blocks and final RMSNorm.

Data

ModelClass

API

class nemo_automodel.components.models.flint.model.FlintForCausalLM(
    config: nemo_automodel.components.models.flint.config.FlintConfig,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)

Bases: HFCheckpointingMixin, PreTrainedModel, MoEFSDPSyncMixin

Eager Flint baseline, with protected trainable sinks and untied FP32 head.

CPU FP32 correctness is the initial supported execution contract. GPU precision, inference caches, packed sequences and distributed axes require their own validation before being advertised.

_keep_in_fp32_modules

= ['_fp32_params', 'lm_head']

_keep_in_fp32_modules_strict

= ['_fp32_params', 'lm_head']

_no_split_modules

= ['FlintDecoderLayer']

_packed_cp_attn_backends

= ('te', 'magi')

_pp_keep_self_forward

bool = True

_pp_return_hidden_states_supported

bool = True

backend

base_model_prefix

= 'model'

cp_mesh

DeviceMesh | None = None

generation_config

= generation_config_from_model_config(config)

lm_head

model

= FlintModel(config, self.backend)

moe_config

packed_mask_type

= 'document_ids'

parallelizer

FlintModelParallelizer = FlintModelParallelizer()

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

vocab_size

= config.vocab_size

nemo_automodel.components.models.flint.model.FlintForCausalLM._is_pipeline_parallel_stage() -> bool

Identify a real partial stage using retained modules and active depth.

nemo_automodel.components.models.flint.model.FlintForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    use_cache: bool = False,
    past_key_values: transformers.cache_utils.Cache | None = None,
    output_hidden_states: bool = False,
    logits_to_keep: int = 0,
    skip_logits: bool = False,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast | torch.Tensor

Return FP32 logits for full-sequence, padded causal language modeling.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Integer [batch, sequence], THD [tokens], or floating CP-local [batch, sequence, hidden] residuals on later PP stages.

attention_mask

torch.Tensor | NoneDefaults to None

Binary [batch, sequence] validity, or nonnegative integer document IDs for eager/SDPA packed attention.

position_ids

torch.Tensor | NoneDefaults to None

Optional integer tensor [batch, sequence].

inputs_embeds

torch.Tensor | NoneDefaults to None

Alternative floating tensor [batch, sequence, hidden].

labels

torch.Tensor | NoneDefaults to None

Optional unshifted HF-style labels [batch, sequence], with -100 ignored. This path shifts once. Production calculate_loss consumes already shifted labels separately and must not pass them here.

use_cache

boolDefaults to False

True is unsupported by this full-sequence training baseline.

past_key_values

Cache | NoneDefaults to None

Inference cache; only None is supported.

output_hidden_states

boolDefaults to False

Include all [batch, sequence, hidden] states, except logits_to_keep=1 selects the final FP32 hidden tensor for the shared linear-loss contract without retaining every layer.

logits_to_keep

intDefaults to 0

Zero projects all positions; a positive value projects that many final positions. Labels require all positions.

skip_logits

boolDefaults to False

Return only final FP32 hidden states for an external already-shifted linear loss; do not run the vocabulary projection.

padding_mask

torch.Tensor | NoneDefaults to None

Boolean [batch, sequence] or THD [tokens], True for physically padded expert rows; document lengths own its meaning.

**attn_kwargs

AnyDefaults to {}

Shared TE THD boundaries/formats or a model-owned Magi plan. PP stage-zero global IDs are selected by prepared local indices once; later stages already receive local residuals.

Returns: CausalLMOutputWithPast | torch.Tensor

CausalLMOutputWithPast with FP32 logits [batch, selected_sequence, vocab],

nemo_automodel.components.models.flint.model.FlintForCausalLM.from_config(
    config: nemo_automodel.components.models.flint.config.FlintConfig,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
) -> nemo_automodel.components.models.flint.model.FlintForCausalLM

classmethod

Construct the native model through the config-owned build boundary.

Parameters:

config

FlintConfig

Local architecture settings.

backend

BackendConfig | NoneDefaults to None

Runtime kernel choices.

Returns: FlintForCausalLM

Initialized native model, without loading any checkpoint weights.

nemo_automodel.components.models.flint.model.FlintForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str | os.PathLike[str],
    model_args: object = (),
    kwargs: object = {}
) -> nemo_automodel.components.models.flint.model.FlintForCausalLM

classmethod

Reject the inherited HF loader, which cannot convert native MoE layouts.

Parameters:

pretrained_model_name_or_path

str | PathLike[str]

Source checkpoint directory or Hub ID.

*model_args

objectDefaults to ()

Upstream loader positional arguments.

**kwargs

objectDefaults to {}

Upstream loader arguments.

Returns: FlintForCausalLM

No model; use NeMoAutoModelForCausalLM.from_pretrained instead.

Raises:

ValueError: Always, with the adapter-aware public entry point.

nemo_automodel.components.models.flint.model.FlintForCausalLM.get_input_embeddings() -> torch.nn.Embedding

Return the independently trainable input embedding module.

Returns: nn.Embedding

Embedding with weight [vocab, hidden], with no frozen padding row.

nemo_automodel.components.models.flint.model.FlintForCausalLM.get_output_embeddings() -> torch.nn.Linear

Return the separate FP32 vocabulary projection module.

Returns: nn.Linear

Bias-free linear module with weight [vocab, hidden].

nemo_automodel.components.models.flint.model.FlintForCausalLM.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]

Describe PP token/residual/logit boundaries without a probe forward.

Parameters:

is_first

bool

Whether this is the first physical stage.

microbatch_size

int

Number of rows in one PP microbatch.

seq_len

int

Global sequence extent, including dataset physical padding.

dtype

torch.dtype

Framework-selected bulk compute dtype.

Returns: tuple[torch.Tensor, ...]

Input/output meta tuples: global int64 IDs [B,S] on stage zero,

nemo_automodel.components.models.flint.model.FlintForCausalLM.initialize_weights(
    dtype: torch.dtype | None = None
) -> None

Initialize every trainable family after construction or meta materialization.

Parameters:

dtype

torch.dtype | NoneDefaults to None

Optional bulk parameter dtype; sinks and the head stay FP32. Installed PEFT adapters retain their explicitly selected dtype.

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

Prepare global [B,S] tokens and [B,N] document lengths before PP splitting.

Parameters:

batch

dict[str, Any]

Global token, target and position tensors; packed batches include real/padded document lengths and qkv_format=‘thd’.

num_chunks

intDefaults to 1

Independent PP microbatch token spaces.

Returns: dict[str, Any]

Model-owned shared sharder and padding metadata; local coordinates

nemo_automodel.components.models.flint.model.FlintForCausalLM.resize_token_embeddings(
    new_num_tokens: int | None = None,
    pad_to_multiple_of: int | None = None,
    mean_resizing: bool = True
) -> torch.nn.Embedding

Preserve the fixed vocabulary used by checkpoint conversion.

Parameters:

new_num_tokens

int | NoneDefaults to None

Vocabulary resizing is not supported by this baseline.

pad_to_multiple_of

int | NoneDefaults to None

Vocabulary padding is not supported.

mean_resizing

boolDefaults to True

Upstream initialization option, unused without resizing.

Returns: nn.Embedding

Existing embedding [vocab, hidden] when no resize is requested.

Raises:

ValueError: A vocabulary resize or padding operation is requested.

nemo_automodel.components.models.flint.model.FlintForCausalLM.set_input_embeddings(
    value: torch.nn.Embedding
) -> None

Replace the input module without changing the output-head parameter.

Parameters:

value

nn.Embedding

Embedding module with weight [vocab, hidden].

nemo_automodel.components.models.flint.model.FlintForCausalLM.set_output_embeddings(
    value: torch.nn.Linear
) -> None

Replace the output module under the protected-FP32 contract.

Parameters:

value

nn.Linear

Bias-free linear module with FP32 weight [vocab, hidden].

class nemo_automodel.components.models.flint.model.FlintModel(
    config: nemo_automodel.components.models.flint.config.FlintConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

Untied input embedding, scheduled decoder blocks and final RMSNorm.

_flint_preserve_output_dtype

bool = True

embed_tokens

hidden_size

= config.hidden_size

layers

norm

residual_dtype

= config.residual_dtype

nemo_automodel.components.models.flint.model.FlintModel.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    output_hidden_states: bool = False,
    attn_kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

Evaluate full-sequence decoder states without an inference cache.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Integer tensor of shape [batch, sequence], or None when inputs_embeds is supplied.

inputs_embeds

torch.Tensor | NoneDefaults to None

Alternative [batch, sequence, hidden] tensor or THD [tokens, hidden]; later PP stages receive this local residual.

position_ids

torch.Tensor | NoneDefaults to None

Optional integer [batch, sequence] or THD [tokens].

attention_mask

torch.Tensor | NoneDefaults to None

Optional binary tensor of shape [batch, sequence]; one means valid, independently of whether a token equals EOS.

padding_mask

torch.Tensor | NoneDefaults to None

Optional True-for-padding expert mask [batch, sequence] or THD [tokens], derived from physical document lengths.

output_hidden_states

boolDefaults to False

Retain embedding, intermediate and final states.

**attn_kwargs

AnyDefaults to {}

Backend THD boundaries/formats or immutable Magi plan.

Returns: BaseModelOutputWithPast

Output with last_hidden_state [batch, sequence, hidden], optional

nemo_automodel.components.models.flint.model.ModelClass = FlintForCausalLM