> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/models/flint/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # nemo_automodel.components.models.flint Flint architecture configuration and native causal language model. ## Submodules * **[`nemo_automodel.components.models.flint.config`](/nemo-automodel/nemo_automodel/components/models/flint/config)** * **[`nemo_automodel.components.models.flint.layers`](/nemo-automodel/nemo_automodel/components/models/flint/layers)** * **[`nemo_automodel.components.models.flint.magi`](/nemo-automodel/nemo_automodel/components/models/flint/magi)** * **[`nemo_automodel.components.models.flint.model`](/nemo-automodel/nemo_automodel/components/models/flint/model)** * **[`nemo_automodel.components.models.flint.packing`](/nemo-automodel/nemo_automodel/components/models/flint/packing)** * **[`nemo_automodel.components.models.flint.parallelization`](/nemo-automodel/nemo_automodel/components/models/flint/parallelization)** * **[`nemo_automodel.components.models.flint.state_dict_adapter`](/nemo-automodel/nemo_automodel/components/models/flint/state_dict_adapter)** ## Package Contents ### Classes | Name | Description | | ------------------------------------------------------------------------------------ | -------------------------------------------------------------------------- | | [`FlintConfig`](#nemo_automodel-components-models-flint-config-FlintConfig) | Architecture settings matching the source Flint checkpoint. | | [`FlintForCausalLM`](#nemo_automodel-components-models-flint-model-FlintForCausalLM) | Eager Flint baseline, with protected trainable sinks and untied FP32 head. | ### API ```python 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` --- ```python nemo_automodel.components.models.flint.config.FlintConfig.__post_init__( kwargs: object = {} ) -> None ``` Resolve schedules and upstream RoPE fields, then validate the contract. ```python 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 | None` — default: None Runtime kernel choices; defaults to the verified eager/torch path. --- **Returns:** `FlintForCausalLM` An initialized native causal language model. ```python nemo_automodel.components.models.flint.config.FlintConfig.validate_layer_type() -> None ``` Validate the preserved source schedule independently of active prefix depth. ```python 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](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `PreTrainedModel`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-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` --- ```python nemo_automodel.components.models.flint.model.FlintForCausalLM._is_pipeline_parallel_stage() -> bool ``` Identify a real partial stage using retained modules and active depth. ```python 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 | None` — default: None Integer \[batch, sequence], THD \[tokens], or floating CP-local \[batch, sequence, hidden] residuals on later PP stages. --- **`attention_mask`** `torch.Tensor | None` — default: None Binary \[batch, sequence] validity, or nonnegative integer document IDs for eager/SDPA packed attention. --- **`position_ids`** `torch.Tensor | None` — default: None Optional integer tensor \[batch, sequence]. --- **`inputs_embeds`** `torch.Tensor | None` — default: None Alternative floating tensor \[batch, sequence, hidden]. --- **`labels`** `torch.Tensor | None` — default: 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`** `bool` — default: False True is unsupported by this full-sequence training baseline. --- **`past_key_values`** `Cache | None` — default: None Inference cache; only None is supported. --- **`output_hidden_states`** `bool` — default: 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`** `int` — default: 0 Zero projects all positions; a positive value projects that many final positions. Labels require all positions. --- **`skip_logits`** `bool` — default: False Return only final FP32 hidden states for an external already-shifted linear loss; do not run the vocabulary projection. --- **`padding_mask`** `torch.Tensor | None` — default: None Boolean \[batch, sequence] or THD \[tokens], True for physically padded expert rows; document lengths own its meaning. --- **`**attn_kwargs`** `Any` — default: \{} 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], ```python 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 | None` — default: None Runtime kernel choices. --- **Returns:** `FlintForCausalLM` Initialized native model, without loading any checkpoint weights. ```python 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`** `object` — default: () Upstream loader positional arguments. --- **`**kwargs`** `object` — default: \{} Upstream loader arguments. --- **Returns:** `FlintForCausalLM` No model; use NeMoAutoModelForCausalLM.from_pretrained instead. **Raises:** * `ValueError`: Always, with the adapter-aware public entry point. ```python 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. ```python 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]. ```python 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, ```python 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 | None` — default: None Optional bulk parameter dtype; sinks and the head stay FP32. Installed PEFT adapters retain their explicitly selected dtype. --- ```python 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`** `int` — default: 1 Independent PP microbatch token spaces. --- **Returns:** `dict[str, Any]` Model-owned shared sharder and padding metadata; local coordinates ```python 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 | None` — default: None Vocabulary resizing is not supported by this baseline. --- **`pad_to_multiple_of`** `int | None` — default: None Vocabulary padding is not supported. --- **`mean_resizing`** `bool` — default: 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. ```python 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]. --- ```python 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]. --- ## Docs - [nemo_automodel.components.models.flint.config](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/models/flint/config.md) - [nemo_automodel.components.models.flint.layers](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/models/flint/layers.md) - [nemo_automodel.components.models.flint.magi](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/models/flint/magi.md) - [nemo_automodel.components.models.flint.model](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/models/flint/model.md) - [nemo_automodel.components.models.flint.packing](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/models/flint/packing.md) - [nemo_automodel.components.models.flint.parallelization](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/models/flint/parallelization.md) - [nemo_automodel.components.models.flint.state_dict_adapter](https://docs.nvidia.com/nemo/automodel/latest/nemo-automodel/nemo_automodel/components/models/flint/state_dict_adapter.md)