nemo_automodel.components.models.flint

Flint architecture configuration and native causal language model.

Submodules

Package Contents

Classes

Name	Description
`FlintConfig`	Architecture settings matching the source Flint checkpoint.
`FlintForCausalLM`	Eager Flint baseline, with protected trainable sinks and untied FP32 head.

API

class nemo_automodel.components.models.flint.config.FlintConfig()

Bases: PreTrainedConfig

Architecture settings matching the source Flint checkpoint.

The local class removes dependence on checkpoint remote code. Full layers have no RoPE; sliding layers rotate all head channels with NeoX RoPE. The window counts the current token. router_aux_loss_coef is preserved as checkpoint metadata; the native training implementation requires an explicit zero value until the original auxiliary objective is specified.

attention_compute_dtype='fp16' is an explicit TE-only QKV/kernel-output boundary within otherwise FP32 projection compute. It does not change the bulk parameter dtype or enable autocast. Small-gradient training requires an appropriate loss-scaling policy; the default None retains native precision.

attention_bias

bool = False

attention_compute_dtype

str | None = None

attention_dropout

float | int = 0.0

bos_token_id

int | None = None

decoder_sparse_step

int = 1

eos_token_id

int | list[int] | None = None

head_dim

int = 128

hidden_act

str = 'silu'

hidden_size

int = 2048

initializer_range

float = 0.02

intermediate_size

int = 6144

keys_to_ignore_at_inference

= ['past_key_values']

layer_types

list[str] | None = None

max_position_embeddings

int = 262144

mlp_only_layers

list[int] | None = None

model_type

= 'flint'

moe_intermediate_size

int = 512

no_rope_layer_indices

list[int] | None = None

norm_topk_prob

bool = True

num_attention_heads

int = 32

num_experts

int = 256

num_experts_per_tok

int = 6

num_hidden_layers

int = 49

num_key_value_heads

int = 4

output_router_logits

bool = False

pad_token_id

int | None = None

residual_dtype

str | None = None

rms_norm_eps

float = 1e-06

rope_parameters

RopeParameters | dict[str, object] | None = None

router_aux_loss_coef

float = 0.0

router_score_function

str = 'sigmoid'

schedule_num_hidden_layers

int | None = None

shared_expert_intermediate_size

int = 1024

sliding_window

int | None = 513

softmax_type

str = 'learnable'

tie_word_embeddings

bool = False

use_cache

bool = False

use_sliding_window

bool = True

vocab_size

int = 150144

nemo_automodel.components.models.flint.config.FlintConfig.__post_init__(
    kwargs: object = {}
) -> None

Resolve schedules and upstream RoPE fields, then validate the contract.

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

Construct a native model without changing these declarative settings.

Parameters:

backend

BackendConfig | NoneDefaults to None

Runtime kernel choices; defaults to the verified eager/torch path.

Returns: FlintForCausalLM

An initialized native causal language model.

nemo_automodel.components.models.flint.config.FlintConfig.validate_layer_type() -> None

Validate the preserved source schedule independently of active prefix depth.

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].