optimum-neuron/pr_1119/en/_app/immutable/nodes/6.bf207ce3.js

import{s as rt,n as jt,o as dt}from"../chunks/scheduler.56725da7.js";import{S as ut,i as ct,e as a,s as t,c as U,h as It,a as y,d as M,b as n,f as mt,g as p,j as i,k as wt,l as ft,m as s,n as T,t as o,o as J,p as C}from"../chunks/index.18a26576.js";import{C as At}from"../chunks/CopyLLMTxtMenu.4edc7eaa.js";import{C as w}from"../chunks/CodeBlock.543bbd4c.js";import{H as m}from"../chunks/MermaidChart.svelte_svelte_type_style_lang.59c1e4a7.js";function ht(Js){let r,me,Je,we,j,re,d,je,u,Cs="This guide explains how to add custom model implementations to the <code>optimum/neuron/models/training/</code> directory. Custom models are needed to support distributed training features like tensor parallelism, pipeline parallelism, and sequence parallelism on AWS Trainium devices.",de,c,ue,I,ce,f,ms="The <code>NeuronModelMixin</code> class provides core functionality:",Ie,A,ws="<li><code>from_pretrained()</code>: Loads regular Transformers weights into custom implementations</li> <li><code>save_pretrained()</code>: Saves sharded checkpoints with consolidation metadata</li> <li>Pipeline parallelism support through <code>PIPELINE_*</code> attributes</li>",fe,h,Ae,b,rs="Transformation specs handle converting weights between:",he,Z,js="<li>Original Transformers format → Custom parallel format (during loading)</li> <li>Custom parallel format → Original Transformers format (during checkpoint consolidation)</li>",be,_,ds="Key transformation spec types:",Ze,W,us="<li><code>FusedLinearsSpec</code>: Handles fused linear layers (e.g., <code>gate_up_proj</code>)</li> <li><code>GQAQKVColumnParallelLinearSpec</code>: Handles grouped query attention projections when the tensor parallel size is greater than the number of key-value heads</li>",_e,g,cs='For complete API documentation of all transformation specs and utility functions, see the <a href="../training_api/transformations">Model Weight Transformation Specs API Reference</a>.',We,V,ge,B,Is="Use these parallel layers from <code>neuronx_distributed</code>:",Ve,k,fs="<li><code>ColumnParallelLinear</code>: Splits weight matrix along output dimension</li> <li><code>RowParallelLinear</code>: Splits weight matrix along input dimension</li> <li><code>ParallelEmbedding</code>: Splits embedding table across ranks</li> <li><code>GQAQKVColumnParallelLinear</code>: Specialized for grouped query attention projections when the tensor parallel size is greater than the number of key-value heads</li>",Be,G,ke,X,Ge,v,As="Create a new directory: <code>optimum/neuron/models/training/your_model/</code>",Xe,$,hs="<strong><code>__init__.py</code></strong>",ve,R,$e,Q,Re,Y,bs="<strong><code>modeling_your_model.py</code></strong>",Qe,F,Ye,N,Fe,x,Ne,z,xe,E,ze,S,Zs="<strong>Important</strong>: Any module that has transformation specs must inherit from <code>CustomModule</code> to ensure proper handling of weight transformations, and the transformation specs must be defined in the <code>self.specs</code> attribute.",Ee,L,Se,H,Le,q,_s="The attention layer implementation depends on the model’s architecture and tensor parallel configuration. There are three main variants:",He,P,Ws="<strong>1. Separate Q, K, V Projections (Default)</strong>",qe,K,Pe,D,gs="<strong>2. Fused QKV Projection (Multi-Head Attention)</strong>",Ke,O,De,ll,Vs="<strong>3. GQA QKV Projection (Required for Challenging TP Configurations)</strong>",Oe,el,lM,Ml,Bs="<strong>When to Use Each Variant:</strong>",eM,sl,ks="<li><strong>Separate Q, K, V</strong>: Default approach, works for all configurations but may be less efficient</li> <li><strong>Fused QKV</strong>: Use when <code>num_heads == num_key_value_heads</code> (no grouped query attention) and <code>fuse_qkv=True</code></li> <li><strong>GQA QKV</strong>: Required when using grouped query attention with challenging tensor parallel configurations where KV heads cannot be evenly distributed across TP ranks</li>",MM,tl,Gs="The choice is typically determined by:",sM,nl,tM,al,nM,yl,aM,il,yM,Ul,iM,pl,UM,Tl,pM,ol,Xs="Update <code>optimum/neuron/models/training/__init__.py</code>:",TM,Jl,oM,Cl,vs="Update <code>optimum/neuron/models/training/auto_models.py</code>:",JM,ml,CM,wl,$s="Here <code>&quot;your_model&quot;</code> is the corresponding to the <code>model_type</code> attribute of your model’s configuration class.",mM,rl,wM,jl,rM,dl,Rs="<li>Use <code>gather_output=False</code> for intermediate layers</li> <li>Set <code>input_is_parallel=True</code> for layers that receive parallel input</li> <li>Configure <code>sequence_parallel_enabled</code> consistently across layers</li> <li>Use appropriate <code>stride</code> values for proper weight sharding</li>",jM,ul,dM,cl,Qs="<li>Always define specs for modules that use fused or parallel layers</li> <li>Use <code>CustomModule</code> mixin for any module with transformation specs</li> <li>Ensure spec parameter names match the actual module structure</li> <li>Test both regular and LoRA weight transformations</li>",uM,Il,cM,fl,Ys="<li>Set <code>SUPPORTS_PIPELINE_PARALLELISM = True</code> for supported models</li> <li>Define <code>PIPELINE_TRANSFORMER_LAYER_CLS</code> as your decoder layer class</li> <li>List all input names in <code>PIPELINE_INPUT_NAMES</code></li>",IM,Al,fM,hl,Fs="<li>Set <code>_supports_flash_attn_2 = True</code> if your model supports it</li> <li>Implement both eager and flash attention paths</li> <li>Use appropriate attention function dispatching</li>",AM,bl,hM,Zl,Ns=`The training tests in <code>tests/training/</code> provide a comprehensive testing framework that validates numerical correctness, distributed training scenarios, and checkpoint compatibility.
Most of the tests are not designed to be run on every custom modeling implementation, but rather to validate the core functionality of the Optimum Neuron training infrastructure.
With that in mind, here’s what you need to implement for your custom modeling:`,bM,_l,ZM,Wl,xs="The <code>test_custom_modeling.py</code> file validates that your custom implementation produces identical outputs to the original Transformers model:",_M,gl,zs="Update <code>tests/training/test_custom_modeling.py</code>:",WM,Vl,gM,Bl,Es="<strong>Important</strong>: For custom modeling validation tests, use small/tiny models to ensure CI efficiency. The test models should have:",VM,kl,Ss="<li>Small vocabulary size (e.g., 1000-8000 tokens)</li> <li>Few layers (e.g., 2-4 layers)</li> <li>Small hidden dimensions (e.g., 128-512)</li> <li>Minimal attention heads (e.g., 4-8 heads)</li>",BM,Gl,Ls="Examples of good test models for custom modeling validation:",kM,Xl,Hs="<li><code>&quot;michaelbenayoun/llama-2-tiny-4kv-heads-4layers-random&quot;</code> - 4 layers, 4 KV heads</li> <li><code>&quot;michaelbenayoun/granite-tiny-4kv-heads-4layers-random&quot;</code> - Tiny Granite model</li> <li><code>&quot;michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random&quot;</code> - Tiny Qwen3 model</li>",GM,vl,qs="Key test your model must pass:",XM,$l,vM,Rl,Ps="<li><strong>Numerical Correctness</strong>: Ensures custom models match Transformers outputs exactly</li> <li><strong>Parallelization Support</strong>: Tests various QKV implementations (regular, fused, GQA)</li>",$M,Ql,RM,Yl,Ks="The <code>test_overfit.py</code> file validates training convergence. To include your model in end-to-end training validation, you must add it to the parametrized test cases:",QM,Fl,Ds="Update <code>tests/training/test_overfit.py</code>:",YM,Nl,FM,xl,Os="This validates:",NM,zl,lt="<li><strong>Convergence Validation</strong>: Ensures models can overfit simple datasets</li>",xM,El,et="Your model will be tested in:",zM,Sl,EM,Ll,SM,Hl,Mt="The <code>test_modeling_auto.py</code> file validates that your model can be loaded using the <code>NeuronModel</code> and <code>NeuronModelForCausalLM</code> auto classes. To include your model in these tests, you must add it to the test cases:",LM,ql,st="Update <code>tests/training/test_modeling_auto.py</code>:",HM,Pl,qM,Kl,tt="This validates:",PM,Dl,nt="<li><strong>Auto Model Loading</strong>: Tests that <code>NeuronModel.from_pretrained()</code> and <code>NeuronModel.from_config()</code> work correctly</li> <li><strong>Auto CausalLM Loading</strong>: Tests that <code>NeuronModelForCausalLM.from_pretrained()</code> and <code>NeuronModelForCausalLM.from_config()</code> work correctly</li>",KM,Ol,DM,le,at="Tests require AWS Trainium instances. Run specific test categories:",OM,ee,ls,Me,es,se,yt="Your implementation must:",Ms,te,it="<li><strong>Pass numerical correctness tests</strong> against original Transformers implementation</li> <li><strong>Support parallelization strategies</strong> (at minimum DP and TP; PP support recommended)</li> <li><strong>Handle various QKV implementations</strong> (regular, fused, GQA)</li> <li><strong>Support checkpoint consolidation</strong> for distributed training</li> <li><strong>Support LoRA training</strong> if applicable</li> <li><strong>Demonstrate convergence</strong> through overfitting tests</li>",ss,ne,Ut="The testing framework ensures your custom model maintains compatibility with the existing Optimum Neuron training infrastructure while delivering expected performance and correctness guarantees.",ts,ae,ns,ye,pt="<li><strong>Weight Shape Mismatches</strong>: Ensure transformation specs handle tensor shapes correctly</li> <li><strong>Pipeline Parallelism Errors</strong>: Check that all required attributes are set</li> <li><strong>Memory Issues</strong>: Consider gradient checkpointing and activation recomputation</li> <li><strong>Attention Compatibility</strong>: Verify attention implementations work with your model architecture</li>",as,ie,ys,Ue,Tt="This guide provides the foundation for implementing custom models. For complete examples and advanced patterns, reference these existing implementations:",is,pe,ot="<li><strong>LLaMA</strong>: <code>optimum/neuron/models/training/llama/modeling_llama.py</code> - Complete implementation with (regular, fused and GQA attention), fused MLP</li> <li><strong>Qwen3</strong>: <code>optimum/neuron/models/training/qwen3/modeling_qwen3.py</code> - Demonstrates how to adapt the Llama implementation for Qwen3 with <code>q_norm</code> and <code>k_norm</code> layers</li>",Us,Te,Jt="Key files to study:",ps,oe,Ct="<li><code>optimum/neuron/models/training/modeling_utils.py</code> - Base <code>NeuronModelMixin</code> class</li> <li><code>optimum/neuron/models/training/transformations_utils.py</code> - Weight transformation specifications</li> <li><code>optimum/neuron/models/training/config.py</code> - <code>TrainingNeuronConfig</code> for parallelism settings</li>",Ts,Ce,os;return j=new At({props:{containerStyle:"float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"}}),d=new m({props:{title:"Contributing Custom Models for Training",local:"contributing-custom-models-for-training",headingTag:"h1"}}),c=new m({props:{title:"Architecture Components",local:"architecture-components",headingTag:"h2"}}),I=new m({props:{title:"1. NeuronModelMixin",local:"1-neuronmodelmixin",headingTag:"h3"}}),h=new m({props:{title:"2. Weight Transformation Specs",local:"2-weight-transformation-specs",headingTag:"h3"}}),V=new m({props:{title:"3. Parallel Layers",local:"3-parallel-layers",headingTag:"h3"}}),G=new m({props:{title:"Implementation Steps",local:"implementation-steps",headingTag:"h2"}}),X=new m({props:{title:"Step 1: Create Model Structure",local:"step-1-create-model-structure",headingTag:"h3"}}),R=new w({props:{code:"ZnJvbSUyMC5tb2RlbGluZ195b3VyX21vZGVsJTIwaW1wb3J0JTIwWW91ck1vZGVsRm9yQ2F1c2FsTE0lMkMlMjBZb3VyTW9kZWwlMEElMEFfX2FsbF9fJTIwJTNEJTIwJTVCJTIyWW91ck1vZGVsRm9yQ2F1c2FsTE0lMjIlMkMlMjAlMjJZb3VyTW9kZWwlMjIlNUQ=",highlighted:`<span class="hljs-keyword">from</span> .modeling_your_model <span class="hljs-keyword">import</span> YourModelForCausalLM, YourModel

__all__ = [<span class="hljs-string">&quot;YourModelForCausalLM&quot;</span>, <span class="hljs-string">&quot;YourModel&quot;</span>]`,lang:"python",wrap:!1}}),Q=new m({props:{title:"Step 2: Implement the Model Building Blocks",local:"step-2-implement-the-model-building-blocks",headingTag:"h3"}}),F=new m({props:{title:"Imports and Dependencies",local:"imports-and-dependencies",headingTag:"h4"}}),N=new w({props:{code:"aW1wb3J0JTIwdG9yY2glMEFmcm9tJTIwdG9yY2glMjBpbXBvcnQlMjBubiUwQWZyb20lMjBuZXVyb254X2Rpc3RyaWJ1dGVkLnBhcmFsbGVsX2xheWVycy5sYXllcnMlMjBpbXBvcnQlMjAoJTBBJTIwJTIwJTIwJTIwQ29sdW1uUGFyYWxsZWxMaW5lYXIlMkMlMEElMjAlMjAlMjAlMjBSb3dQYXJhbGxlbExpbmVhciUyQyUwQSUyMCUyMCUyMCUyMFBhcmFsbGVsRW1iZWRkaW5nJTJDJTBBKSUwQWZyb20lMjBuZXVyb254X2Rpc3RyaWJ1dGVkLm1vZHVsZXMucWt2X2xpbmVhciUyMGltcG9ydCUyMEdRQVFLVkNvbHVtblBhcmFsbGVsTGluZWFyJTBBZnJvbSUyMHRyYW5zZm9ybWVycyUyMGltcG9ydCUyMFByZVRyYWluZWRNb2RlbCUwQWZyb20lMjB0cmFuc2Zvcm1lcnMubW9kZWxzLnlvdXJfbW9kZWwlMjBpbXBvcnQlMjBZb3VyTW9kZWxDb25maWclMEElMEFmcm9tJTIwLi5jb25maWclMjBpbXBvcnQlMjBUcmFpbmluZ05ldXJvbkNvbmZpZyUwQWZyb20lMjAuLm1vZGVsaW5nX3V0aWxzJTIwaW1wb3J0JTIwTmV1cm9uTW9kZWxNaXhpbiUwQWZyb20lMjAuLnRyYW5zZm9ybWF0aW9uc191dGlscyUyMGltcG9ydCUyMCglMEElMjAlMjAlMjAlMjBDdXN0b21Nb2R1bGUlMkMlMEElMjAlMjAlMjAlMjBGdXNlZExpbmVhcnNTcGVjJTJDJTBBJTIwJTIwJTIwJTIwR1FBUUtWQ29sdW1uUGFyYWxsZWxMaW5lYXJTcGVjJTJDJTBBJTIwJTIwJTIwJTIwTW9kZWxXZWlnaHRUcmFuc2Zvcm1hdGlvblNwZWNzJTJDJTBBKQ==",highlighted:`<span class="hljs-keyword">import</span> torch
<span class="hljs-keyword">from</span> torch <span class="hljs-keyword">import</span> nn
<span class="hljs-keyword">from</span> neuronx_distributed.parallel_layers.layers <span class="hljs-keyword">import</span> (
    ColumnParallelLinear,
    RowParallelLinear,
    ParallelEmbedding,
)
<span class="hljs-keyword">from</span> neuronx_distributed.modules.qkv_linear <span class="hljs-keyword">import</span> GQAQKVColumnParallelLinear
<span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> PreTrainedModel
<span class="hljs-keyword">from</span> transformers.models.your_model <span class="hljs-keyword">import</span> YourModelConfig

<span class="hljs-keyword">from</span> ..config <span class="hljs-keyword">import</span> TrainingNeuronConfig
<span class="hljs-keyword">from</span> ..modeling_utils <span class="hljs-keyword">import</span> NeuronModelMixin
<span class="hljs-keyword">from</span> ..transformations_utils <span class="hljs-keyword">import</span> (
    CustomModule,
    FusedLinearsSpec,
    GQAQKVColumnParallelLinearSpec,
    ModelWeightTransformationSpecs,
)`,lang:"python",wrap:!1}}),x=new m({props:{title:"Embedding Layer",local:"embedding-layer",headingTag:"h4"}}),z=new w({props:{code:"Y2xhc3MlMjBZb3VyTW9kZWxFbWJlZGRpbmdzKG5uLk1vZHVsZSklM0ElMEElMjAlMjAlMjAlMjBkZWYlMjBfX2luaXRfXyhzZWxmJTJDJTIwY29uZmlnJTJDJTIwdHJuX2NvbmZpZyklM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzdXBlcigpLl9faW5pdF9fKCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmVtYmVkX3Rva2VucyUyMCUzRCUyMFBhcmFsbGVsRW1iZWRkaW5nKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGNvbmZpZy52b2NhYl9zaXplJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwY29uZmlnLmhpZGRlbl9zaXplJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZHR5cGUlM0Rjb25maWcuZHR5cGUlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZXF1ZW5jZV9wYXJhbGxlbF9lbmFibGVkJTNEdHJuX2NvbmZpZy5zZXF1ZW5jZV9wYXJhbGxlbF9lbmFibGVkJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwKQ==",highlighted:`<span class="hljs-keyword">class</span> <span class="hljs-title class_">YourModelEmbeddings</span>(nn.Module):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, config, trn_config</span>):
        <span class="hljs-built_in">super</span>().__init__()
        self.embed_tokens = ParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
            dtype=config.dtype,
            sequence_parallel_enabled=trn_config.sequence_parallel_enabled,
        )`,lang:"python",wrap:!1}}),E=new m({props:{title:"MLP Layer with Fused Linears",local:"mlp-layer-with-fused-linears",headingTag:"h4"}}),L=new w({props:{code:"Y2xhc3MlMjBZb3VyTW9kZWxNTFAobm4uTW9kdWxlJTJDJTIwQ3VzdG9tTW9kdWxlKSUzQSUwQSUyMCUyMCUyMCUyMGRlZiUyMF9faW5pdF9fKHNlbGYlMkMlMjBjb25maWclMkMlMjB0cm5fY29uZmlnKSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHN1cGVyKCkuX19pbml0X18oKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYuaGlkZGVuX3NpemUlMjAlM0QlMjBjb25maWcuaGlkZGVuX3NpemUlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmludGVybWVkaWF0ZV9zaXplJTIwJTNEJTIwY29uZmlnLmludGVybWVkaWF0ZV9zaXplJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIzJTIwRnVzZWQlMjBnYXRlJTIwYW5kJTIwdXAlMjBwcm9qZWN0aW9ucyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYuZ2F0ZV91cF9wcm9qJTIwJTNEJTIwQ29sdW1uUGFyYWxsZWxMaW5lYXIoJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5oaWRkZW5fc2l6ZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMDIlMjAqJTIwc2VsZi5pbnRlcm1lZGlhdGVfc2l6ZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHN0cmlkZSUzRDIlMkMlMjAlMjAlMjMlMjBJbXBvcnRhbnQlMjBmb3IlMjBwcm9wZXIlMjBzaGFyZGluZyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGJpYXMlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGdhdGhlcl9vdXRwdXQlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlcXVlbmNlX3BhcmFsbGVsX2VuYWJsZWQlM0R0cm5fY29uZmlnLnNlcXVlbmNlX3BhcmFsbGVsX2VuYWJsZWQlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBkdHlwZSUzRGNvbmZpZy5kdHlwZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmRvd25fcHJvaiUyMCUzRCUyMFJvd1BhcmFsbGVsTGluZWFyKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYuaW50ZXJtZWRpYXRlX3NpemUlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmhpZGRlbl9zaXplJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwYmlhcyUzREZhbHNlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwaW5wdXRfaXNfcGFyYWxsZWwlM0RUcnVlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VxdWVuY2VfcGFyYWxsZWxfZW5hYmxlZCUzRHRybl9jb25maWcuc2VxdWVuY2VfcGFyYWxsZWxfZW5hYmxlZCUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGR0eXBlJTNEY29uZmlnLmR0eXBlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyMERlZmluZSUyMHRyYW5zZm9ybWF0aW9uJTIwc3BlY3MlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLnNwZWNzJTIwJTNEJTIwTW9kZWxXZWlnaHRUcmFuc2Zvcm1hdGlvblNwZWNzKCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLnNwZWNzLmFkZF9zcGVjKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMEZ1c2VkTGluZWFyc1NwZWMoJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZnVzZWRfbGluZWFyX25hbWUlM0QlMjJnYXRlX3VwX3Byb2olMjIlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBsaW5lYXJfbmFtZXMlM0QlNUIlMjJnYXRlX3Byb2olMjIlMkMlMjAlMjJ1cF9wcm9qJTIyJTVEJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwYmlhcyUzREZhbHNlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZnVzZV9heGlzJTNEJTIyY29sdW1uJTIyJTJDJTIwJTIwJTIzJTIwRnVzZSUyMGFsb25nJTIwb3V0cHV0JTIwZGltZW5zaW9uJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwb3JpZ2luYWxfZGltcyUzRCU1QnNlbGYuaW50ZXJtZWRpYXRlX3NpemUlMkMlMjBzZWxmLmludGVybWVkaWF0ZV9zaXplJTVEJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCk=",highlighted:`<span class="hljs-keyword">class</span> <span class="hljs-title class_">YourModelMLP</span>(nn.Module, CustomModule):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, config, trn_config</span>):
        <span class="hljs-built_in">super</span>().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        
        <span class="hljs-comment"># Fused gate and up projections</span>
        self.gate_up_proj = ColumnParallelLinear(
            self.hidden_size,
            <span class="hljs-number">2</span> * self.intermediate_size,
            stride=<span class="hljs-number">2</span>,  <span class="hljs-comment"># Important for proper sharding</span>
            bias=<span class="hljs-literal">False</span>,
            gather_output=<span class="hljs-literal">False</span>,
            sequence_parallel_enabled=trn_config.sequence_parallel_enabled,
            dtype=config.dtype,
        )
        
        self.down_proj = RowParallelLinear(
            self.intermediate_size,
            self.hidden_size,
            bias=<span class="hljs-literal">False</span>,
            input_is_parallel=<span class="hljs-literal">True</span>,
            sequence_parallel_enabled=trn_config.sequence_parallel_enabled,
            dtype=config.dtype,
        )
        
        <span class="hljs-comment"># Define transformation specs</span>
        self.specs = ModelWeightTransformationSpecs()
        self.specs.add_spec(
            FusedLinearsSpec(
                fused_linear_name=<span class="hljs-string">&quot;gate_up_proj&quot;</span>,
                linear_names=[<span class="hljs-string">&quot;gate_proj&quot;</span>, <span class="hljs-string">&quot;up_proj&quot;</span>],
                bias=<span class="hljs-literal">False</span>,
                fuse_axis=<span class="hljs-string">&quot;column&quot;</span>,  <span class="hljs-comment"># Fuse along output dimension</span>
                original_dims=[self.intermediate_size, self.intermediate_size],
            )
        )`,lang:"python",wrap:!1}}),H=new m({props:{title:"Attention Layer",local:"attention-layer",headingTag:"h4"}}),K=new w({props:{code:"Y2xhc3MlMjBZb3VyTW9kZWxBdHRlbnRpb24obm4uTW9kdWxlJTJDJTIwQ3VzdG9tTW9kdWxlKSUzQSUwQSUyMCUyMCUyMCUyMGRlZiUyMF9faW5pdF9fKHNlbGYlMkMlMjBjb25maWclMkMlMjB0cm5fY29uZmlnJTJDJTIwbGF5ZXJfaWR4KSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHN1cGVyKCkuX19pbml0X18oKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYuY29uZmlnJTIwJTNEJTIwY29uZmlnJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5udW1faGVhZHMlMjAlM0QlMjBjb25maWcubnVtX2F0dGVudGlvbl9oZWFkcyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYubnVtX2tleV92YWx1ZV9oZWFkcyUyMCUzRCUyMGNvbmZpZy5udW1fa2V5X3ZhbHVlX2hlYWRzJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5oZWFkX2RpbSUyMCUzRCUyMGNvbmZpZy5oaWRkZW5fc2l6ZSUyMCUyRiUyRiUyMHNlbGYubnVtX2hlYWRzJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIzJTIwU2VwYXJhdGUlMjBwcm9qZWN0aW9ucyUyMGZvciUyMFElMkMlMjBLJTJDJTIwViUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYucV9wcm9qJTIwJTNEJTIwQ29sdW1uUGFyYWxsZWxMaW5lYXIoJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwY29uZmlnLmhpZGRlbl9zaXplJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5udW1faGVhZHMlMjAqJTIwc2VsZi5oZWFkX2RpbSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGJpYXMlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGdhdGhlcl9vdXRwdXQlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlcXVlbmNlX3BhcmFsbGVsX2VuYWJsZWQlM0R0cm5fY29uZmlnLnNlcXVlbmNlX3BhcmFsbGVsX2VuYWJsZWQlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBkdHlwZSUzRGNvbmZpZy5kdHlwZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmtfcHJvaiUyMCUzRCUyMENvbHVtblBhcmFsbGVsTGluZWFyKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGNvbmZpZy5oaWRkZW5fc2l6ZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYubnVtX2tleV92YWx1ZV9oZWFkcyUyMColMjBzZWxmLmhlYWRfZGltJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwYmlhcyUzREZhbHNlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZ2F0aGVyX291dHB1dCUzREZhbHNlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VxdWVuY2VfcGFyYWxsZWxfZW5hYmxlZCUzRHRybl9jb25maWcuc2VxdWVuY2VfcGFyYWxsZWxfZW5hYmxlZCUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGR0eXBlJTNEY29uZmlnLmR0eXBlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYudl9wcm9qJTIwJTNEJTIwQ29sdW1uUGFyYWxsZWxMaW5lYXIoJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwY29uZmlnLmhpZGRlbl9zaXplJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5udW1fa2V5X3ZhbHVlX2hlYWRzJTIwKiUyMHNlbGYuaGVhZF9kaW0lMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBiaWFzJTNERmFsc2UlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBnYXRoZXJfb3V0cHV0JTNERmFsc2UlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZXF1ZW5jZV9wYXJhbGxlbF9lbmFibGVkJTNEdHJuX2NvbmZpZy5zZXF1ZW5jZV9wYXJhbGxlbF9lbmFibGVkJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZHR5cGUlM0Rjb25maWcuZHR5cGUlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjApJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5vX3Byb2olMjAlM0QlMjBSb3dQYXJhbGxlbExpbmVhciglMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLm51bV9oZWFkcyUyMColMjBzZWxmLmhlYWRfZGltJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwY29uZmlnLmhpZGRlbl9zaXplJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwYmlhcyUzREZhbHNlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwaW5wdXRfaXNfcGFyYWxsZWwlM0RUcnVlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VxdWVuY2VfcGFyYWxsZWxfZW5hYmxlZCUzRHRybl9jb25maWcuc2VxdWVuY2VfcGFyYWxsZWxfZW5hYmxlZCUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGR0eXBlJTNEY29uZmlnLmR0eXBlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyME5vJTIwdHJhbnNmb3JtYXRpb24lMjBzcGVjcyUyMG5lZWRlZCUyMC0lMjByZWd1bGFyJTIwcGFyYWxsZWwlMjBsYXllcnMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLnNwZWNzJTIwJTNEJTIwTW9kZWxXZWlnaHRUcmFuc2Zvcm1hdGlvblNwZWNzKCk=",highlighted:`<span class="hljs-keyword">class</span> <span class="hljs-title class_">YourModelAttention</span>(nn.Module, CustomModule):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, config, trn_config, layer_idx</span>):
        <span class="hljs-built_in">super</span>().__init__()
        self.config = config
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // self.num_heads
        
        <span class="hljs-comment"># Separate projections for Q, K, V</span>
        self.q_proj = ColumnParallelLinear(
            config.hidden_size,
            self.num_heads * self.head_dim,
            bias=<span class="hljs-literal">False</span>,
            gather_output=<span class="hljs-literal">False</span>,
            sequence_parallel_enabled=trn_config.sequence_parallel_enabled,
            dtype=config.dtype,
        )
        self.k_proj = ColumnParallelLinear(
            config.hidden_size,
            self.num_key_value_heads * self.head_dim,
            bias=<span class="hljs-literal">False</span>,
            gather_output=<span class="hljs-literal">False</span>,
            sequence_parallel_enabled=trn_config.sequence_parallel_enabled,
            dtype=config.dtype,
        )
        self.v_proj = ColumnParallelLinear(
            config.hidden_size,
            self.num_key_value_heads * self.head_dim,
            bias=<span class="hljs-literal">False</span>,
            gather_output=<span class="hljs-literal">False</span>,
            sequence_parallel_enabled=trn_config.sequence_parallel_enabled,
            dtype=config.dtype,
        )
        
        self.o_proj = RowParallelLinear(
            self.num_heads * self.head_dim,
            config.hidden_size,
            bias=<span class="hljs-literal">False</span>,
            input_is_parallel=<span class="hljs-literal">True</span>,
            sequence_parallel_enabled=trn_config.sequence_parallel_enabled,
            dtype=config.dtype,
        )
        
        <span class="hljs-comment"># No transformation specs needed - regular parallel layers</span>
        self.specs = ModelWeightTransformationSpecs()`,lang:"python",wrap:!1}}),O=new w({props:{code:"Y2xhc3MlMjBZb3VyTW9kZWxBdHRlbnRpb24obm4uTW9kdWxlJTJDJTIwQ3VzdG9tTW9kdWxlKSUzQSUwQSUyMCUyMCUyMCUyMGRlZiUyMF9faW5pdF9fKHNlbGYlMkMlMjBjb25maWclMkMlMjB0cm5fY29uZmlnJTJDJTIwbGF5ZXJfaWR4KSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHN1cGVyKCkuX19pbml0X18oKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyMC4uLiUyMChzYW1lJTIwc2V0dXAlMjBhcyUyMGFib3ZlKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHRwX3NpemUlMjAlM0QlMjBnZXRfdGVuc29yX21vZGVsX3BhcmFsbGVsX3NpemUoKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyME9ubHklMjB1c2UlMjBmdXNlZCUyMFFLViUyMHdoZW4lMjBudW1faGVhZHMlMjAlM0QlM0QlMjBudW1fa2V5X3ZhbHVlX2hlYWRzJTIwKG5vJTIwR1FBKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGlmJTIwdHJuX2NvbmZpZy5mdXNlX3FrdiUyMGFuZCUyMHNlbGYubnVtX2hlYWRzJTIwJTNEJTNEJTIwc2VsZi5udW1fa2V5X3ZhbHVlX2hlYWRzJTNBJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5xa3ZfcHJvaiUyMCUzRCUyMENvbHVtblBhcmFsbGVsTGluZWFyKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGNvbmZpZy5oaWRkZW5fc2l6ZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMDMlMjAqJTIwc2VsZi5udW1faGVhZHMlMjAqJTIwc2VsZi5oZWFkX2RpbSUyQyUyMCUyMCUyMyUyMFElMjAlMkIlMjBLJTIwJTJCJTIwViUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHN0cmlkZSUzRDMlMkMlMjAlMjAlMjMlMjBJbXBvcnRhbnQlMjBmb3IlMjBwcm9wZXIlMjBzaGFyZGluZyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGJpYXMlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGdhdGhlcl9vdXRwdXQlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlcXVlbmNlX3BhcmFsbGVsX2VuYWJsZWQlM0R0cm5fY29uZmlnLnNlcXVlbmNlX3BhcmFsbGVsX2VuYWJsZWQlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBkdHlwZSUzRGNvbmZpZy5kdHlwZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjMlMjBEZWZpbmUlMjB0cmFuc2Zvcm1hdGlvbiUyMHNwZWNzJTIwZm9yJTIwZnVzZWQlMjBRS1YlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLnNwZWNzJTIwJTNEJTIwTW9kZWxXZWlnaHRUcmFuc2Zvcm1hdGlvblNwZWNzKCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLnNwZWNzLmFkZF9zcGVjKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMEZ1c2VkTGluZWFyc1NwZWMoJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZnVzZWRfbGluZWFyX25hbWUlM0QlMjJxa3ZfcHJvaiUyMiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGxpbmVhcl9uYW1lcyUzRCU1QiUyMnFfcHJvaiUyMiUyQyUyMCUyMmtfcHJvaiUyMiUyQyUyMCUyMnZfcHJvaiUyMiU1RCUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGJpYXMlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGZ1c2VfYXhpcyUzRCUyMmNvbHVtbiUyMiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMG9yaWdpbmFsX2RpbXMlM0QlNUJzZWxmLm51bV9oZWFkcyUyMColMjBzZWxmLmhlYWRfZGltJTVEJTIwKiUyMDMlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjApJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYuc3BsaXRfc2l6ZSUyMCUzRCUyMHNlbGYubnVtX2hlYWRzJTIwKiUyMHNlbGYuaGVhZF9kaW0lMjAlMkYlMkYlMjB0cF9zaXpl",highlighted:`<span class="hljs-keyword">class</span> <span class="hljs-title class_">YourModelAttention</span>(nn.Module, CustomModule):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, config, trn_config, layer_idx</span>):
        <span class="hljs-built_in">super</span>().__init__()
        <span class="hljs-comment"># ... (same setup as above)</span>
        
        tp_size = get_tensor_model_parallel_size()
        
        <span class="hljs-comment"># Only use fused QKV when num_heads == num_key_value_heads (no GQA)</span>
        <span class="hljs-keyword">if</span> trn_config.fuse_qkv <span class="hljs-keyword">and</span> self.num_heads == self.num_key_value_heads:
            self.qkv_proj = ColumnParallelLinear(
                config.hidden_size,
                <span class="hljs-number">3</span> * self.num_heads * self.head_dim,  <span class="hljs-comment"># Q + K + V</span>
                stride=<span class="hljs-number">3</span>,  <span class="hljs-comment"># Important for proper sharding</span>
                bias=<span class="hljs-literal">False</span>,
                gather_output=<span class="hljs-literal">False</span>,
                sequence_parallel_enabled=trn_config.sequence_parallel_enabled,
                dtype=config.dtype,
            )
            
            <span class="hljs-comment"># Define transformation specs for fused QKV</span>
            self.specs = ModelWeightTransformationSpecs()
            self.specs.add_spec(
                FusedLinearsSpec(
                    fused_linear_name=<span class="hljs-string">&quot;qkv_proj&quot;</span>,
                    linear_names=[<span class="hljs-string">&quot;q_proj&quot;</span>, <span class="hljs-string">&quot;k_proj&quot;</span>, <span class="hljs-string">&quot;v_proj&quot;</span>],
                    bias=<span class="hljs-literal">False</span>,
                    fuse_axis=<span class="hljs-string">&quot;column&quot;</span>,
                    original_dims=[self.num_heads * self.head_dim] * <span class="hljs-number">3</span>,
                )
            )
            self.split_size = self.num_heads * self.head_dim // tp_size`,lang:"python",wrap:!1}}),el=new w({props:{code:"Y2xhc3MlMjBZb3VyTW9kZWxBdHRlbnRpb24obm4uTW9kdWxlJTJDJTIwQ3VzdG9tTW9kdWxlKSUzQSUwQSUyMCUyMCUyMCUyMGRlZiUyMF9faW5pdF9fKHNlbGYlMkMlMjBjb25maWclMkMlMjB0cm5fY29uZmlnJTJDJTIwbGF5ZXJfaWR4KSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHN1cGVyKCkuX19pbml0X18oKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyMC4uLiUyMChzYW1lJTIwc2V0dXAlMjBhcyUyMGFib3ZlKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHRwX3NpemUlMjAlM0QlMjBnZXRfdGVuc29yX21vZGVsX3BhcmFsbGVsX3NpemUoKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyMFVzZSUyMEdRQSUyMFFLViUyMHdoZW4lMjBLViUyMGhlYWRzJTIwY2FuJ3QlMjBiZSUyMGV2ZW5seSUyMGRpc3RyaWJ1dGVkJTIwYWNyb3NzJTIwVFAlMjByYW5rcyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyMFRoaXMlMjBoYXBwZW5zJTIwd2hlbiUzQSUyMG51bV9rZXlfdmFsdWVfaGVhZHMlMjAlM0MlMjB0cF9zaXplJTIwb3IlMjBudW1fa2V5X3ZhbHVlX2hlYWRzJTIwJTI1JTIwdHBfc2l6ZSUyMCElM0QlMjAwJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5xa3ZfbGluZWFyJTIwJTNEJTIwKHNlbGYubnVtX2tleV92YWx1ZV9oZWFkcyUyMCUzQyUyMHRwX3NpemUpJTIwb3IlMjAoc2VsZi5udW1fa2V5X3ZhbHVlX2hlYWRzJTIwJTI1JTIwdHBfc2l6ZSUyMCElM0QlMjAwKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGlmJTIwc2VsZi5xa3ZfbGluZWFyJTNBJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIzJTIwQ2FsY3VsYXRlJTIwS1YlMjBzaXplJTIwbXVsdGlwbGllciUyMHRvJTIwZW5zdXJlJTIwZXZlbiUyMGRpc3RyaWJ1dGlvbiUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGlmJTIwdHJuX2NvbmZpZy5rdl9zaXplX211bHRpcGxpZXIlMjBpcyUyME5vbmUlM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmt2X3NpemVfbXVsdGlwbGllciUyMCUzRCUyMHRybl9jb25maWcuYXV0b19rdl9zaXplX211bHRpcGxpZXIoc2VsZi5udW1fa2V5X3ZhbHVlX2hlYWRzKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGVsc2UlM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmt2X3NpemVfbXVsdGlwbGllciUyMCUzRCUyMHRybl9jb25maWcua3Zfc2l6ZV9tdWx0aXBsaWVyJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5xa3ZfcHJvaiUyMCUzRCUyMEdRQVFLVkNvbHVtblBhcmFsbGVsTGluZWFyKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGNvbmZpZy5oaWRkZW5fc2l6ZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCU1QnNlbGYubnVtX2hlYWRzJTIwKiUyMHNlbGYuaGVhZF9kaW0lMkMlMjBzZWxmLm51bV9rZXlfdmFsdWVfaGVhZHMlMjAqJTIwc2VsZi5oZWFkX2RpbSU1RCUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGJpYXMlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGdhdGhlcl9vdXRwdXQlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlcXVlbmNlX3BhcmFsbGVsX2VuYWJsZWQlM0R0cm5fY29uZmlnLnNlcXVlbmNlX3BhcmFsbGVsX2VuYWJsZWQlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBrdl9zaXplX211bHRpcGxpZXIlM0RzZWxmLmt2X3NpemVfbXVsdGlwbGllciUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGZ1c2VfcWt2JTNEdHJuX2NvbmZpZy5mdXNlX3FrdiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGR0eXBlJTNEY29uZmlnLmR0eXBlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyMERlZmluZSUyMHRyYW5zZm9ybWF0aW9uJTIwc3BlY3MlMjBmb3IlMjBHUUElMjBRS1YlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLnNwZWNzJTIwJTNEJTIwTW9kZWxXZWlnaHRUcmFuc2Zvcm1hdGlvblNwZWNzKCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLnNwZWNzLmFkZF9zcGVjKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMEdRQVFLVkNvbHVtblBhcmFsbGVsTGluZWFyU3BlYyglMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBncWFfcWt2X3Byb2plY3Rpb25fbmFtZSUzRCUyMnFrdl9wcm9qJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwcXVlcnlfcHJvamVjdGlvbl9uYW1lJTNEJTIycV9wcm9qJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwa2V5X3Byb2plY3Rpb25fbmFtZSUzRCUyMmtfcHJvaiUyMiUyQyUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHZhbHVlX3Byb2plY3Rpb25fbmFtZSUzRCUyMnZfcHJvaiUyMiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMG91dHB1dF9wcm9qZWN0aW9uX25hbWUlM0QlMjJvX3Byb2olMjIlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBudW1fYXR0ZW50aW9uX2hlYWRzJTNEc2VsZi5udW1faGVhZHMlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBudW1fa2V5X3ZhbHVlX2hlYWRzJTNEc2VsZi5udW1fa2V5X3ZhbHVlX2hlYWRzJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwa3Zfc2l6ZV9tdWx0aXBsaWVyJTNEc2VsZi5rdl9zaXplX211bHRpcGxpZXIlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBxX291dHB1dF9zaXplX3Blcl9wYXJ0aXRpb24lM0RzZWxmLnFrdl9wcm9qLnFfb3V0cHV0X3NpemVfcGVyX3BhcnRpdGlvbiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGt2X291dHB1dF9zaXplX3Blcl9wYXJ0aXRpb24lM0RzZWxmLnFrdl9wcm9qLmt2X291dHB1dF9zaXplX3Blcl9wYXJ0aXRpb24lMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBmdXNlX3FrdiUzRHRybl9jb25maWcuZnVzZV9xa3YlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjApJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwKQ==",highlighted:`<span class="hljs-keyword">class</span> <span class="hljs-title class_">YourModelAttention</span>(nn.Module, CustomModule):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, config, trn_config, layer_idx</span>):
        <span class="hljs-built_in">super</span>().__init__()
        <span class="hljs-comment"># ... (same setup as above)</span>
        
        tp_size = get_tensor_model_parallel_size()
        
        <span class="hljs-comment"># Use GQA QKV when KV heads can&#x27;t be evenly distributed across TP ranks</span>
        <span class="hljs-comment"># This happens when: num_key_value_heads &lt; tp_size or num_key_value_heads % tp_size != 0</span>
        self.qkv_linear = (self.num_key_value_heads &lt; tp_size) <span class="hljs-keyword">or</span> (self.num_key_value_heads % tp_size != <span class="hljs-number">0</span>)
        
        <span class="hljs-keyword">if</span> self.qkv_linear:
            <span class="hljs-comment"># Calculate KV size multiplier to ensure even distribution</span>
            <span class="hljs-keyword">if</span> trn_config.kv_size_multiplier <span class="hljs-keyword">is</span> <span class="hljs-literal">None</span>:
                self.kv_size_multiplier = trn_config.auto_kv_size_multiplier(self.num_key_value_heads)
            <span class="hljs-keyword">else</span>:
                self.kv_size_multiplier = trn_config.kv_size_multiplier
                
            self.qkv_proj = GQAQKVColumnParallelLinear(
                config.hidden_size,
                [self.num_heads * self.head_dim, self.num_key_value_heads * self.head_dim],
                bias=<span class="hljs-literal">False</span>,
                gather_output=<span class="hljs-literal">False</span>,
                sequence_parallel_enabled=trn_config.sequence_parallel_enabled,
                kv_size_multiplier=self.kv_size_multiplier,
                fuse_qkv=trn_config.fuse_qkv,
                dtype=config.dtype,
            )
            
            <span class="hljs-comment"># Define transformation specs for GQA QKV</span>
            self.specs = ModelWeightTransformationSpecs()
            self.specs.add_spec(
                GQAQKVColumnParallelLinearSpec(
                    gqa_qkv_projection_name=<span class="hljs-string">&quot;qkv_proj&quot;</span>,
                    query_projection_name=<span class="hljs-string">&quot;q_proj&quot;</span>,
                    key_projection_name=<span class="hljs-string">&quot;k_proj&quot;</span>, 
                    value_projection_name=<span class="hljs-string">&quot;v_proj&quot;</span>,
                    output_projection_name=<span class="hljs-string">&quot;o_proj&quot;</span>,
                    num_attention_heads=self.num_heads,
                    num_key_value_heads=self.num_key_value_heads,
                    kv_size_multiplier=self.kv_size_multiplier,
                    q_output_size_per_partition=self.qkv_proj.q_output_size_per_partition,
                    kv_output_size_per_partition=self.qkv_proj.kv_output_size_per_partition,
                    fuse_qkv=trn_config.fuse_qkv,
                )
            )`,lang:"python",wrap:!1}}),nl=new w({props:{code:"dHBfc2l6ZSUyMCUzRCUyMGdldF90ZW5zb3JfbW9kZWxfcGFyYWxsZWxfc2l6ZSgpJTBBdXNlX2dxYV9xa3YlMjAlM0QlMjAobnVtX2tleV92YWx1ZV9oZWFkcyUyMCUzQyUyMHRwX3NpemUpJTIwb3IlMjAobnVtX2tleV92YWx1ZV9oZWFkcyUyMCUyNSUyMHRwX3NpemUlMjAhJTNEJTIwMCklMEF1c2VfZnVzZWRfcWt2JTIwJTNEJTIwdHJuX2NvbmZpZy5mdXNlX3FrdiUyMGFuZCUyMChudW1faGVhZHMlMjAlM0QlM0QlMjBudW1fa2V5X3ZhbHVlX2hlYWRzKSUyMGFuZCUyMG5vdCUyMHVzZV9ncWFfcWt2",highlighted:`tp_size = get_tensor_model_parallel_size()
use_gqa_qkv = (num_key_value_heads &lt; tp_size) <span class="hljs-keyword">or</span> (num_key_value_heads % tp_size != <span class="hljs-number">0</span>)
use_fused_qkv = trn_config.fuse_qkv <span class="hljs-keyword">and</span> (num_heads == num_key_value_heads) <span class="hljs-keyword">and</span> <span class="hljs-keyword">not</span> use_gqa_qkv`,lang:"python",wrap:!1}}),al=new m({props:{title:"Step 3: Implement Main Model Classes",local:"step-3-implement-main-model-classes",headingTag:"h3"}}),yl=new m({props:{title:"Base Model",local:"base-model",headingTag:"h4"}}),il=new w({props:{code:"Y2xhc3MlMjBZb3VyUHJlVHJhaW5lZE1vZGVsKFByZVRyYWluZWRNb2RlbCUyQyUyME5ldXJvbk1vZGVsTWl4aW4pJTNBJTBBJTIwJTIwJTIwJTIwY29uZmlnX2NsYXNzJTIwJTNEJTIwWW91ck1vZGVsQ29uZmlnJTBBJTIwJTIwJTIwJTIwYmFzZV9tb2RlbF9wcmVmaXglMjAlM0QlMjAlMjJtb2RlbCUyMiUwQSUyMCUyMCUyMCUyMHN1cHBvcnRzX2dyYWRpZW50X2NoZWNrcG9pbnRpbmclMjAlM0QlMjBUcnVlJTBBJTIwJTIwJTIwJTIwX25vX3NwbGl0X21vZHVsZXMlMjAlM0QlMjAlNUIlMjJZb3VyTW9kZWxEZWNvZGVyTGF5ZXIlMjIlNUQlMEElMjAlMjAlMjAlMjBfc2tpcF9rZXlzX2RldmljZV9wbGFjZW1lbnQlMjAlM0QlMjAlMjJwYXN0X2tleV92YWx1ZXMlMjIlMEElMjAlMjAlMjAlMjBfc3VwcG9ydHNfZmxhc2hfYXR0bl8yJTIwJTNEJTIwVHJ1ZSUwQSUyMCUyMCUyMCUyMF9zdXBwb3J0c19jYWNoZV9jbGFzcyUyMCUzRCUyMFRydWUlMEElMjAlMjAlMjAlMjBfc3VwcG9ydHNfcXVhbnRpemVkX2NhY2hlJTIwJTNEJTIwVHJ1ZSUwQSUyMCUyMCUyMCUyMF9zdXBwb3J0c19zdGF0aWNfY2FjaGUlMjAlM0QlMjBUcnVlJTBBJTBBJTBBY2xhc3MlMjBZb3VyTW9kZWwoTmV1cm9uTW9kZWxNaXhpbiUyQyUyMFlvdXJQcmVUcmFpbmVkTW9kZWwpJTNBJTBBJTIwJTIwJTIwJTIwZGVmJTIwX19pbml0X18oc2VsZiUyQyUyMGNvbmZpZyUzQSUyMFlvdXJNb2RlbENvbmZpZyUyQyUyMHRybl9jb25maWclM0ElMjBUcmFpbmluZ05ldXJvbkNvbmZpZyklM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBZb3VyUHJlVHJhaW5lZE1vZGVsLl9faW5pdF9fKHNlbGYlMkMlMjBjb25maWcpJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi5wYWRkaW5nX2lkeCUyMCUzRCUyMGNvbmZpZy5wYWRfdG9rZW5faWQlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLnZvY2FiX3NpemUlMjAlM0QlMjBjb25maWcudm9jYWJfc2l6ZSUwQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYudHJuX2NvbmZpZyUyMCUzRCUyMHRybl9jb25maWclMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmVtYmVkX3Rva2VucyUyMCUzRCUyMFBhcmFsbGVsRW1iZWRkaW5nKC4uLiklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmxheWVycyUyMCUzRCUyMG5uLk1vZHVsZUxpc3QoJTVCJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwWW91ck1vZGVsRGVjb2RlckxheWVyKGNvbmZpZyUyQyUyMHRybl9jb25maWclMkMlMjBsYXllcl9pZHgpJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwZm9yJTIwbGF5ZXJfaWR4JTIwaW4lMjByYW5nZShjb25maWcubnVtX2hpZGRlbl9sYXllcnMpJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTVEKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYubm9ybSUyMCUzRCUyMFlvdXJNb2RlbFJNU05vcm0oLi4uKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYucG9zdF9pbml0KCk=",highlighted:`<span class="hljs-keyword">class</span> <span class="hljs-title class_">YourPreTrainedModel</span>(PreTrainedModel, NeuronModelMixin):
    config_class = YourModelConfig
    base_model_prefix = <span class="hljs-string">&quot;model&quot;</span>
    supports_gradient_checkpointing = <span class="hljs-literal">True</span>
    _no_split_modules = [<span class="hljs-string">&quot;YourModelDecoderLayer&quot;</span>]
    _skip_keys_device_placement = <span class="hljs-string">&quot;past_key_values&quot;</span>
    _supports_flash_attn_2 = <span class="hljs-literal">True</span>
    _supports_cache_class = <span class="hljs-literal">True</span>
    _supports_quantized_cache = <span class="hljs-literal">True</span>
    _supports_static_cache = <span class="hljs-literal">True</span>


<span class="hljs-keyword">class</span> <span class="hljs-title class_">YourModel</span>(NeuronModelMixin, YourPreTrainedModel):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, config: YourModelConfig, trn_config: TrainingNeuronConfig</span>):
        YourPreTrainedModel.__init__(self, config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.trn_config = trn_config
        
        self.embed_tokens = ParallelEmbedding(...)
        self.layers = nn.ModuleList([
            YourModelDecoderLayer(config, trn_config, layer_idx)
            <span class="hljs-keyword">for</span> layer_idx <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(config.num_hidden_layers)
        ])
        self.norm = YourModelRMSNorm(...)
        
        self.post_init()`,lang:"python",wrap:!1}}),Ul=new m({props:{title:"CausalLM Model",local:"causallm-model",headingTag:"h4"}}),pl=new w({props:{code:"Y2xhc3MlMjBZb3VyTW9kZWxGb3JDYXVzYWxMTShOZXVyb25Nb2RlbE1peGluJTJDJTIwWW91clByZVRyYWluZWRNb2RlbCklM0ElMEElMjAlMjAlMjAlMjBfdGllZF93ZWlnaHRzX2tleXMlMjAlM0QlMjAlNUIlMjJsbV9oZWFkLndlaWdodCUyMiU1RCUwQSUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMyUyMFBpcGVsaW5lJTIwcGFyYWxsZWxpc20lMjBzdXBwb3J0JTBBJTIwJTIwJTIwJTIwU1VQUE9SVFNfUElQRUxJTkVfUEFSQUxMRUxJU00lMjAlM0QlMjBUcnVlJTBBJTIwJTIwJTIwJTIwUElQRUxJTkVfVFJBTlNGT1JNRVJfTEFZRVJfQ0xTJTIwJTNEJTIwWW91ck1vZGVsRGVjb2RlckxheWVyJTBBJTIwJTIwJTIwJTIwUElQRUxJTkVfSU5QVVRfTkFNRVMlMjAlM0QlMjAlNUIlMjJpbnB1dF9pZHMlMjIlMkMlMjAlMjJhdHRlbnRpb25fbWFzayUyMiU1RCUwQSUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMGRlZiUyMF9faW5pdF9fKHNlbGYlMkMlMjBjb25maWclMkMlMjB0cm5fY29uZmlnKSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHN1cGVyKCkuX19pbml0X18oY29uZmlnKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYudHJuX2NvbmZpZyUyMCUzRCUyMHRybl9jb25maWclMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLm1vZGVsJTIwJTNEJTIwWW91ck1vZGVsKGNvbmZpZyUyQyUyMHRybl9jb25maWcpJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwc2VsZi52b2NhYl9zaXplJTIwJTNEJTIwY29uZmlnLnZvY2FiX3NpemUlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzZWxmLmxtX2hlYWQlMjAlM0QlMjBDb2x1bW5QYXJhbGxlbExpbmVhciglMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBjb25maWcuaGlkZGVuX3NpemUlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBjb25maWcudm9jYWJfc2l6ZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGJpYXMlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGdhdGhlcl9vdXRwdXQlM0RGYWxzZSUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGR0eXBlJTNEY29uZmlnLmR0eXBlJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMHNlbGYucG9zdF9pbml0KCk=",highlighted:`<span class="hljs-keyword">class</span> <span class="hljs-title class_">YourModelForCausalLM</span>(NeuronModelMixin, YourPreTrainedModel):
    _tied_weights_keys = [<span class="hljs-string">&quot;lm_head.weight&quot;</span>]
    
    <span class="hljs-comment"># Pipeline parallelism support</span>
    SUPPORTS_PIPELINE_PARALLELISM = <span class="hljs-literal">True</span>
    PIPELINE_TRANSFORMER_LAYER_CLS = YourModelDecoderLayer
    PIPELINE_INPUT_NAMES = [<span class="hljs-string">&quot;input_ids&quot;</span>, <span class="hljs-string">&quot;attention_mask&quot;</span>]
    
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, config, trn_config</span>):
        <span class="hljs-built_in">super</span>().__init__(config)
        self.trn_config = trn_config
        self.model = YourModel(config, trn_config)
        self.vocab_size = config.vocab_size
        
        self.lm_head = ColumnParallelLinear(
            config.hidden_size,
            config.vocab_size,
            bias=<span class="hljs-literal">False</span>,
            gather_output=<span class="hljs-literal">False</span>,
            dtype=config.dtype,
        )
        
        self.post_init()`,lang:"python",wrap:!1}}),Tl=new m({props:{title:"Step 4: Register Model",local:"step-4-register-model",headingTag:"h3"}}),Jl=new w({props:{code:"ZnJvbSUyMC55b3VyX21vZGVsJTIwaW1wb3J0JTIwWW91ck1vZGVsRm9yQ2F1c2FsTE0lMkMlMjBZb3VyTW9kZWwlMEElMEFfX2FsbF9fJTIwJTNEJTIwJTVCLi4uJTJDJTIwJTIyWW91ck1vZGVsRm9yQ2F1c2FsTE0lMjIlMkMlMjAlMjJZb3VyTW9kZWwlMjIlNUQ=",highlighted:`<span class="hljs-keyword">from</span> .your_model <span class="hljs-keyword">import</span> YourModelForCausalLM, YourModel

__all__ = [..., <span class="hljs-string">&quot;YourModelForCausalLM&quot;</span>, <span class="hljs-string">&quot;YourModel&quot;</span>]`,lang:"python",wrap:!1}}),ml=new w({props:{code:"ZnJvbSUyMC55b3VyX21vZGVsLm1vZGVsaW5nX3lvdXJfbW9kZWwlMjBpbXBvcnQlMjBZb3VyTW9kZWxGb3JDYXVzYWxMTSUyQyUyMFlvdXJNb2RlbCUwQSUwQSUyMyUyMFJlZ2lzdGVyJTIwdGhlJTIwYmFzZSUyMG1vZGVsJTIwKHdpdGhvdXQlMjBoZWFkKSUwQXJlZ2lzdGVyX25ldXJvbl9tb2RlbF9mb3JfdHJhaW5pbmcoJTIyeW91cl9tb2RlbCUyMiUyQyUyMCUyMm1vZGVsJTIyKShZb3VyTW9kZWwpJTBBJTBBJTIzJTIwUmVnaXN0ZXIlMjB0aGUlMjBDYXVzYWxMTSUyMG1vZGVsJTBBcmVnaXN0ZXJfbmV1cm9uX21vZGVsX2Zvcl90cmFpbmluZyglMjJ5b3VyX21vZGVsJTIyJTJDJTIwJTIydGV4dC1nZW5lcmF0aW9uJTIyKShZb3VyTW9kZWxGb3JDYXVzYWxMTSk=",highlighted:`<span class="hljs-keyword">from</span> .your_model.modeling_your_model <span class="hljs-keyword">import</span> YourModelForCausalLM, YourModel

<span class="hljs-comment"># Register the base model (without head)</span>
register_neuron_model_for_training(<span class="hljs-string">&quot;your_model&quot;</span>, <span class="hljs-string">&quot;model&quot;</span>)(YourModel)

<span class="hljs-comment"># Register the CausalLM model</span>
register_neuron_model_for_training(<span class="hljs-string">&quot;your_model&quot;</span>, <span class="hljs-string">&quot;text-generation&quot;</span>)(YourModelForCausalLM)`,lang:"python",wrap:!1}}),rl=new m({props:{title:"Best Practices",local:"best-practices",headingTag:"h2"}}),jl=new m({props:{title:"1. Parallel Layer Configuration",local:"1-parallel-layer-configuration",headingTag:"h3"}}),ul=new m({props:{title:"2. Weight Transformation Specs",local:"2-weight-transformation-specs",headingTag:"h3"}}),Il=new m({props:{title:"3. Pipeline Parallelism",local:"3-pipeline-parallelism",headingTag:"h3"}}),Al=new m({props:{title:"4. Flash Attention Support",local:"4-flash-attention-support",headingTag:"h3"}}),bl=new m({props:{title:"Testing Your Implementation",local:"testing-your-implementation",headingTag:"h2"}}),_l=new m({props:{title:"1. Custom Modeling Validation",local:"1-custom-modeling-validation",headingTag:"h3"}}),Vl=new w({props:{code:"Q1VTVE9NX01PREVMSU5HU19UT19URVNUJTIwJTNEJTIwJTVCJTBBJTIwJTIwJTIwJTIwJTIzJTIwLi4uJTIwZXhpc3RpbmclMjBtb2RlbHMlMjAuLi4lMEElMjAlMjAlMjAlMjAoJTIyWW91ck1vZGVsRm9yQ2F1c2FsTE0lMjIlMkMlMjAlMjJ5b3VyLW9yZyUyRnlvdXItbW9kZWwtbmFtZSUyMiklMkMlMEElNUQ=",highlighted:`CUSTOM_MODELINGS_TO_TEST = [
    <span class="hljs-comment"># ... existing models ...</span>
    (<span class="hljs-string">&quot;YourModelForCausalLM&quot;</span>, <span class="hljs-string">&quot;your-org/your-model-name&quot;</span>),
]`,lang:"python",wrap:!1}}),$l=new w({props:{code:"ZGVmJTIwdGVzdF9jdXN0b21fbW9kZWxpbmdfbWF0Y2hlc19vcmlnaW5hbCgpJTIwJTIwJTIzJTIwT3V0cHV0JTIwbWF0Y2hpbmc=",highlighted:'<span class="hljs-keyword">def</span> <span class="hljs-title function_">test_custom_modeling_matches_original</span>()  <span class="hljs-comment"># Output matching</span>',lang:"python",wrap:!1}}),Ql=new m({props:{title:"2. End-to-End Training Validation",local:"2-end-to-end-training-validation",headingTag:"h3"}}),Nl=new w({props:{code:"JTQwcHl0ZXN0Lm1hcmsucGFyYW1ldHJpemUoJTBBJTIwJTIwJTIwJTIwJTIybW9kZWxfY2xhc3NfbmFtZSUyQ21vZGVsX25hbWVfb3JfcGF0aCUyQ2xlYXJuaW5nX3JhdGUlMkN3YXJtdXBfcmF0aW8lMkN0cmFpbmluZ19rd2FyZ3MlMkN1c2VfZmxhc2hfYXR0ZW50aW9uXzIlMkNtYXhfZXhwZWN0ZWRfbG9zcyUyQ21heF9sZW5ndGglMkNudW1fc3RlcHMlMjIlMkMlMEElMjAlMjAlMjAlMjAlNUIlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjMlMjAuLi4lMjBleGlzdGluZyUyMG1vZGVscyUyMC4uLiUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCU1QiUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMllvdXJNb2RlbEZvckNhdXNhbExNJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIyeW91ci1vcmclMkZ5b3VyLW1vZGVsLW5hbWUlMjIlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAxZS00JTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwMC4wMyUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCU3QiU3RCUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMFRydWUlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAwLjUlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAyMDQ4JTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwNTAlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlNUQlMkMlMEElMjAlMjAlMjAlMjAlNUQlMkMlMEElMjAlMjAlMjAlMjBpZHMlM0QlNUIlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjMlMjAuLi4lMjBleGlzdGluZyUyMG1vZGVsJTIwSURzJTIwLi4uJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIyeW91ci1vcmclMkZ5b3VyLW1vZGVsLW5hbWUlMjIlMkMlMEElMjAlMjAlMjAlMjAlNUQlMkMlMEEp",highlighted:`<span class="hljs-meta">@pytest.mark.parametrize(<span class="hljs-params">
    <span class="hljs-string">&quot;model_class_name,model_name_or_path,learning_rate,warmup_ratio,training_kwargs,use_flash_attention_2,max_expected_loss,max_length,num_steps&quot;</span>,
    [
        <span class="hljs-comment"># ... existing models ...</span>
        [
            <span class="hljs-string">&quot;YourModelForCausalLM&quot;</span>,
            <span class="hljs-string">&quot;your-org/your-model-name&quot;</span>,
            <span class="hljs-number">1e-4</span>,
            <span class="hljs-number">0.03</span>,
            {},
            <span class="hljs-literal">True</span>,
            <span class="hljs-number">0.5</span>,
            <span class="hljs-number">2048</span>,
            <span class="hljs-number">50</span>,
        ],
    ],
    ids=[
        <span class="hljs-comment"># ... existing model IDs ...</span>
        <span class="hljs-string">&quot;your-org/your-model-name&quot;</span>,
    ],
</span>)</span>`,lang:"python",wrap:!1}}),Sl=new w({props:{code:"ZGVmJTIwdGVzdF9vdmVyZml0X2N1c3RvbV9tb2RlbGluZ19jYXVzYWxfbG0oKSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMyUyMEJhc2ljJTIwdHJhaW5pbmclMjAoeW91ciUyMG1vZGVsJTIwaW5jbHVkZWQp",highlighted:'<span class="hljs-keyword">def</span> <span class="hljs-title function_">test_overfit_custom_modeling_causal_lm</span>()       <span class="hljs-comment"># Basic training (your model included)</span>',lang:"python",wrap:!1}}),Ll=new m({props:{title:"3. Auto Model Loading",local:"3-auto-model-loading",headingTag:"h3"}}),Pl=new w({props:{code:"JTQwcHl0ZXN0Lm1hcmsucGFyYW1ldHJpemUoJTIyZnJvbV9wcmV0cmFpbmVkJTIyJTJDJTIwJTVCRmFsc2UlMkMlMjBUcnVlJTVEJTJDJTIwaWRzJTNEJTVCJTIyZnJvbV9jb25maWclMjIlMkMlMjAlMjJmcm9tX3ByZXRyYWluZWQlMjIlNUQpJTBBJTQwZGlzdHJpYnV0ZWRfdGVzdCh3b3JsZF9zaXplJTNEMSklMEElNDBpc190cmFpbml1bV90ZXN0JTBBZGVmJTIwdGVzdF9hdXRvX21vZGVsX3dpdGhfc3VwcG9ydGVkX2FyY2hpdGVjdHVyZShmcm9tX3ByZXRyYWluZWQpJTNBJTBBJTIwJTIwJTIwJTIwdHJuX2NvbmZpZyUyMCUzRCUyMFRyYWluaW5nTmV1cm9uQ29uZmlnKCklMEElMjAlMjAlMjAlMjBrd2FyZ3MlMjAlM0QlMjAlN0IlMjJkdHlwZSUyMiUzQSUyMHRvcmNoLmJmbG9hdDE2JTdEJTBBJTIwJTIwJTIwJTIwZm9yJTIwbW9kZWxfbmFtZV9vcl9wYXRoJTIwaW4lMjAlNUIlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJtaWNoYWVsYmVuYXlvdW4lMkZsbGFtYS0yLXRpbnktNGt2LWhlYWRzLTRsYXllcnMtcmFuZG9tJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIybWljaGFlbGJlbmF5b3VuJTJGZ3Jhbml0ZS10aW55LTRrdi1oZWFkcy00bGF5ZXJzLXJhbmRvbSUyMiUyQyUyMCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMm1pY2hhZWxiZW5heW91biUyRnF3ZW4zLXRpbnktNGt2LWhlYWRzLTRsYXllcnMtcmFuZG9tJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIyeW91ci1vcmclMkZ5b3VyLW1vZGVsLW5hbWUlMjIlMkMlMjAlMjAlMjMlMjBBZGQlMjB5b3VyJTIwbW9kZWwlMjBoZXJlJTBBJTIwJTIwJTIwJTIwJTVEJTNBJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIzJTIwLi4uJTIwcmVzdCUyMG9mJTIwdGVzdCUyMGxvZ2ljJTBBJTBBJTQwcHl0ZXN0Lm1hcmsucGFyYW1ldHJpemUoJTIyZnJvbV9wcmV0cmFpbmVkJTIyJTJDJTIwJTVCRmFsc2UlMkMlMjBUcnVlJTVEJTJDJTIwaWRzJTNEJTVCJTIyZnJvbV9jb25maWclMjIlMkMlMjAlMjJmcm9tX3ByZXRyYWluZWQlMjIlNUQpJTBBJTQwZGlzdHJpYnV0ZWRfdGVzdCh3b3JsZF9zaXplJTNEMSklMEElNDBpc190cmFpbml1bV90ZXN0JTBBZGVmJTIwdGVzdF9hdXRvX21vZGVsX2Zvcl9jYXVzYWxfbG1fd2l0aF9zdXBwb3J0ZWRfYXJjaGl0ZWN0dXJlKGZyb21fcHJldHJhaW5lZCklM0ElMEElMjAlMjAlMjAlMjB0cm5fY29uZmlnJTIwJTNEJTIwVHJhaW5pbmdOZXVyb25Db25maWcoKSUwQSUyMCUyMCUyMCUyMGt3YXJncyUyMCUzRCUyMCU3QiUyMmR0eXBlJTIyJTNBJTIwdG9yY2guYmZsb2F0MTYlN0QlMEElMjAlMjAlMjAlMjBmb3IlMjBtb2RlbF9uYW1lX29yX3BhdGglMjBpbiUyMCU1QiUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMm1pY2hhZWxiZW5heW91biUyRmxsYW1hLTItdGlueS00a3YtaGVhZHMtNGxheWVycy1yYW5kb20lMjIlMkMlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJtaWNoYWVsYmVuYXlvdW4lMkZncmFuaXRlLXRpbnktNGt2LWhlYWRzLTRsYXllcnMtcmFuZG9tJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIybWljaGFlbGJlbmF5b3VuJTJGcXdlbjMtdGlueS00a3YtaGVhZHMtNGxheWVycy1yYW5kb20lMjIlMkMlMjAlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjJ5b3VyLW9yZyUyRnlvdXItbW9kZWwtbmFtZSUyMiUyQyUyMCUyMCUyMyUyMEFkZCUyMHlvdXIlMjBtb2RlbCUyMGhlcmUlMEElMjAlMjAlMjAlMjAlNUQlM0ElMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjMlMjAuLi4lMjByZXN0JTIwb2YlMjB0ZXN0JTIwbG9naWM=",highlighted:`<span class="hljs-meta">@pytest.mark.parametrize(<span class="hljs-params"><span class="hljs-string">&quot;from_pretrained&quot;</span>, [<span class="hljs-literal">False</span>, <span class="hljs-literal">True</span>], ids=[<span class="hljs-string">&quot;from_config&quot;</span>, <span class="hljs-string">&quot;from_pretrained&quot;</span>]</span>)</span>
<span class="hljs-meta">@distributed_test(<span class="hljs-params">world_size=<span class="hljs-number">1</span></span>)</span>
<span class="hljs-meta">@is_trainium_test</span>
<span class="hljs-keyword">def</span> <span class="hljs-title function_">test_auto_model_with_supported_architecture</span>(<span class="hljs-params">from_pretrained</span>):
    trn_config = TrainingNeuronConfig()
    kwargs = {<span class="hljs-string">&quot;dtype&quot;</span>: torch.bfloat16}
    <span class="hljs-keyword">for</span> model_name_or_path <span class="hljs-keyword">in</span> [
        <span class="hljs-string">&quot;michaelbenayoun/llama-2-tiny-4kv-heads-4layers-random&quot;</span>,
        <span class="hljs-string">&quot;michaelbenayoun/granite-tiny-4kv-heads-4layers-random&quot;</span>, 
        <span class="hljs-string">&quot;michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random&quot;</span>,
        <span class="hljs-string">&quot;your-org/your-model-name&quot;</span>,  <span class="hljs-comment"># Add your model here</span>
    ]:
        <span class="hljs-comment"># ... rest of test logic</span>

<span class="hljs-meta">@pytest.mark.parametrize(<span class="hljs-params"><span class="hljs-string">&quot;from_pretrained&quot;</span>, [<span class="hljs-literal">False</span>, <span class="hljs-literal">True</span>], ids=[<span class="hljs-string">&quot;from_config&quot;</span>, <span class="hljs-string">&quot;from_pretrained&quot;</span>]</span>)</span>
<span class="hljs-meta">@distributed_test(<span class="hljs-params">world_size=<span class="hljs-number">1</span></span>)</span>
<span class="hljs-meta">@is_trainium_test</span>
<span class="hljs-keyword">def</span> <span class="hljs-title function_">test_auto_model_for_causal_lm_with_supported_architecture</span>(<span class="hljs-params">from_pretrained</span>):
    trn_config = TrainingNeuronConfig()
    kwargs = {<span class="hljs-string">&quot;dtype&quot;</span>: torch.bfloat16}
    <span class="hljs-keyword">for</span> model_name_or_path <span class="hljs-keyword">in</span> [
        <span class="hljs-string">&quot;michaelbenayoun/llama-2-tiny-4kv-heads-4layers-random&quot;</span>,
        <span class="hljs-string">&quot;michaelbenayoun/granite-tiny-4kv-heads-4layers-random&quot;</span>,
        <span class="hljs-string">&quot;michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random&quot;</span>, 
        <span class="hljs-string">&quot;your-org/your-model-name&quot;</span>,  <span class="hljs-comment"># Add your model here</span>
    ]:
        <span class="hljs-comment"># ... rest of test logic</span>`,lang:"python",wrap:!1}}),Ol=new m({props:{title:"4. Running Tests",local:"4-running-tests",headingTag:"h3"}}),ee=new w({props:{code:"JTIzJTIwUnVuJTIwYWxsJTIwY3VzdG9tJTIwbW9kZWxpbmclMjB0ZXN0cyUwQXB5dGVzdCUyMHRlc3RzJTJGdHJhaW5pbmclMkZ0ZXN0X2N1c3RvbV9tb2RlbGluZy5weSUyMC12JTBBJTBBJTIzJTIwUnVuJTIwc3BlY2lmaWMlMjBtb2RlbCUyMHRlc3RzJTBBcHl0ZXN0JTIwdGVzdHMlMkZ0cmFpbmluZyUyRnRlc3RfY3VzdG9tX21vZGVsaW5nLnB5JTIwLXYlMjAtayUyMCUyMnlvdXJfbW9kZWwlMjIlMEElMEElMjMlMjBSdW4lMjBlbmQtdG8tZW5kJTIwdHJhaW5pbmclMjB2YWxpZGF0aW9uJTBBcHl0ZXN0JTIwdGVzdHMlMkZ0cmFpbmluZyUyRnRlc3Rfb3ZlcmZpdC5weSUyMC12",highlighted:`<span class="hljs-comment"># Run all custom modeling tests</span>
pytest tests/training/test_custom_modeling.py -v

<span class="hljs-comment"># Run specific model tests</span>
pytest tests/training/test_custom_modeling.py -v -k <span class="hljs-string">&quot;your_model&quot;</span>

<span class="hljs-comment"># Run end-to-end training validation</span>
pytest tests/training/test_overfit.py -v`,lang:"bash",wrap:!1}}),Me=new m({props:{title:"5. Test Requirements",local:"5-test-requirements",headingTag:"h3"}}),ae=new m({props:{title:"Common Issues",local:"common-issues",headingTag:"h2"}}),ie=new m({props:{title:"Additional Resources",local:"additional-resources",headingTag:"h2"}}),{c(){r=a("meta"),me=t(),Je=a("p"),we=t(),U(j.$$.fragment),re=t(),U(d.$$.fragment),je=t(),u=a("p"),u.innerHTML=Cs,de=t(),U(c.$$.fragment),ue=t(),U(I.$$.fragment),ce=t(),f=a("p"),f.innerHTML=ms,Ie=t(),A=a("ul"),A.innerHTML=ws,fe=t(),U(h.$$.fragment),Ae=t(),b=a("p"),b.textContent=rs,he=t(),Z=a("ul"),Z.innerHTML=js,be=t(),_=a("p"),_.textContent=ds,Ze=t(),W=a("ul"),W.innerHTML=us,_e=t(),g=a("p"),g.innerHTML=cs,We=t(),U(V.$$.fragment),ge=t(),B=a("p"),B.innerHTML=Is,Ve=t(),k=a("ul"),k.innerHTML=fs,Be=t(),U(G.$$.fragment),ke=t(),U(X.$$.fragment),Ge=t(),v=a("p"),v.innerHTML=As,Xe=t(),$=a("p"),$.innerHTML=hs,ve=t(),U(R.$$.fragment),$e=t(),U(Q.$$.fragment),Re=t(),Y=a("p"),Y.innerHTML=bs,Qe=t(),U(F.$$.fragment),Ye=t(),U(N.$$.fragment),Fe=t(),U(x.$$.fragment),Ne=t(),U(z.$$.fragment),xe=t(),U(E.$$.fragment),ze=t(),S=a("p"),S.innerHTML=Zs,Ee=t(),U(L.$$.fragment),Se=t(),U(H.$$.fragment),Le=t(),q=a("p"),q.textContent=_s,He=t(),P=a("p"),P.innerHTML=Ws,qe=t(),U(K.$$.fragment),Pe=t(),D=a("p"),D.innerHTML=gs,Ke=t(),U(O.$$.fragment),De=t(),ll=a("p"),ll.innerHTML=Vs,Oe=t(),U(el.$$.fragment),lM=t(),Ml=a("p"),Ml.innerHTML=Bs,eM=t(),sl=a("ul"),sl.innerHTML=ks,MM=t(),tl=a("p"),tl.textContent=Gs,sM=t(),U(nl.$$.fragment),tM=t(),U(al.$$.fragment),nM=t(),U(yl.$$.fragment),aM=t(),U(il.$$.fragment),yM=t(),U(Ul.$$.fragment),iM=t(),U(pl.$$.fragment),UM=t(),U(Tl.$$.fragment),pM=t(),ol=a("p"),ol.innerHTML=Xs,TM=t(),U(Jl.$$.fragment),oM=t(),Cl=a("p"),Cl.innerHTML=vs,JM=t(),U(ml.$$.fragment),CM=t(),wl=a("p"),wl.innerHTML=$s,mM=t(),U(rl.$$.fragment),wM=t(),U(jl.$$.fragment),rM=t(),dl=a("ul"),dl.innerHTML=Rs,jM=t(),U(ul.$$.fragment),dM=t(),cl=a("ul"),cl.innerHTML=Qs,uM=t(),U(Il.$$.fragment),cM=t(),fl=a("ul"),fl.innerHTML=Ys,IM=t(),U(Al.$$.fragment),fM=t(),hl=a("ul"),hl.innerHTML=Fs,AM=t(),U(bl.$$.fragment),hM=t(),Zl=a("p"),Zl.innerHTML=Ns,bM=t(),U(_l.$$.fragment),ZM=t(),Wl=a("p"),Wl.innerHTML=xs,_M=t(),gl=a("p"),gl.innerHTML=zs,WM=t(),U(Vl.$$.fragment),gM=t(),Bl=a("p"),Bl.innerHTML=Es,VM=t(),kl=a("ul"),kl.innerHTML=Ss,BM=t(),Gl=a("p"),Gl.textContent=Ls,kM=t(),Xl=a("ul"),Xl.innerHTML=Hs,GM=t(),vl=a("p"),vl.textContent=qs,XM=t(),U($l.$$.fragment),vM=t(),Rl=a("ul"),Rl.innerHTML=Ps,$M=t(),U(Ql.$$.fragment),RM=t(),Yl=a("p"),Yl.innerHTML=Ks,QM=t(),Fl=a("p"),Fl.innerHTML=Ds,YM=t(),U(Nl.$$.fragment),FM=t(),xl=a("p"),xl.textContent=Os,NM=t(),zl=a("ul"),zl.innerHTML=lt,xM=t(),El=a("p"),El.textContent=et,zM=t(),U(Sl.$$.fragment),EM=t(),U(Ll.$$.fragment),SM=t(),Hl=a("p"),Hl.innerHTML=Mt,LM=t(),ql=a("p"),ql.innerHTML=st,HM=t(),U(Pl.$$.fragment),qM=t(),Kl=a("p"),Kl.textContent=tt,PM=t(),Dl=a("ul"),Dl.innerHTML=nt,KM=t(),U(Ol.$$.fragment),DM=t(),le=a("p"),le.textContent=at,OM=t(),U(ee.$$.fragment),ls=t(),U(Me.$$.fragment),es=t(),se=a("p"),se.textContent=yt,Ms=t(),te=a("ol"),te.innerHTML=it,ss=t(),ne=a("p"),ne.textContent=Ut,ts=t(),U(ae.$$.fragment),ns=t(),ye=a("ul"),ye.innerHTML=pt,as=t(),U(ie.$$.fragment),ys=t(),Ue=a("p"),Ue.textContent=Tt,is=t(),pe=a("ul"),pe.innerHTML=ot,Us=t(),Te=a("p"),Te.textContent=Jt,ps=t(),oe=a("ul"),oe.innerHTML=Ct,Ts=t(),Ce=a("p"),this.h()},l(l){const e=It("svelte-u9bgzb",document.head);r=y(e,"META",{name:!0,content:!0}),e.forEach(M),me=n(l),Je=y(l,"P",{}),mt(Je).forEach(M),we=n(l),p(j.$$.fragment,l),re=n(l),p(d.$$.fragment,l),je=n(l),u=y(l,"P",{"data-svelte-h":!0}),i(u)!=="svelte-o1icnh"&&(u.innerHTML=Cs),de=n(l),p(c.$$.fragment,l),ue=n(l),p(I.$$.fragment,l),ce=n(l),f=y(l,"P",{"data-svelte-h":!0}),i(f)!=="svelte-o8lxl9"&&(f.innerHTML=ms),Ie=n(l),A=y(l,"UL",{"data-svelte-h":!0}),i(A)!=="svelte-17o51ne"&&(A.innerHTML=ws),fe=n(l),p(h.$$.fragment,l),Ae=n(l),b=y(l,"P",{"data-svelte-h":!0}),i(b)!=="svelte-1cxti7r"&&(b.textContent=rs),he=n(l),Z=y(l,"UL",{"data-svelte-h":!0}),i(Z)!=="svelte-egt1g"&&(Z.innerHTML=js),be=n(l),_=y(l,"P",{"data-svelte-h":!0}),i(_)!=="svelte-1sw314o"&&(_.textContent=ds),Ze=n(l),W=y(l,"UL",{"data-svelte-h":!0}),i(W)!=="svelte-1813bw9"&&(W.innerHTML=us),_e=n(l),g=y(l,"P",{"data-svelte-h":!0}),i(g)!=="svelte-1xgxhwy"&&(g.innerHTML=cs),We=n(l),p(V.$$.fragment,l),ge=n(l),B=y(l,"P",{"data-svelte-h":!0}),i(B)!=="svelte-1xa5nbr"&&(B.innerHTML=Is),Ve=n(l),k=y(l,"UL",{"data-svelte-h":!0}),i(k)!=="svelte-1p13ho0"&&(k.innerHTML=fs),Be=n(l),p(G.$$.fragment,l),ke=n(l),p(X.$$.fragment,l),Ge=n(l),v=y(l,"P",{"data-svelte-h":!0}),i(v)!=="svelte-1brmy4r"&&(v.innerHTML=As),Xe=n(l),$=y(l,"P",{"data-svelte-h":!0}),i($)!=="svelte-1bum30f"&&($.innerHTML=hs),ve=n(l),p(R.$$.fragment,l),$e=n(l),p(Q.$$.fragment,l),Re=n(l),Y=y(l,"P",{"data-svelte-h":!0}),i(Y)!=="svelte-1a3kcfq"&&(Y.innerHTML=bs),Qe=n(l),p(F.$$.fragment,l),Ye=n(l),p(N.$$.fragment,l),Fe=n(l),p(x.$$.fragment,l),Ne=n(l),p(z.$$.fragment,l),xe=n(l),p(E.$$.fragment,l),ze=n(l),S=y(l,"P",{"data-svelte-h":!0}),i(S)!=="svelte-1kxohx9"&&(S.innerHTML=Zs),Ee=n(l),p(L.$$.fragment,l),Se=n(l),p(H.$$.fragment,l),Le=n(l),q=y(l,"P",{"data-svelte-h":!0}),i(q)!=="svelte-67vg61"&&(q.textContent=_s),He=n(l),P=y(l,"P",{"data-svelte-h":!0}),i(P)!=="svelte-b4subx"&&(P.innerHTML=Ws),qe=n(l),p(K.$$.fragment,l),Pe=n(l),D=y(l,"P",{"data-svelte-h":!0}),i(D)!=="svelte-f9hwk6"&&(D.innerHTML=gs),Ke=n(l),p(O.$$.fragment,l),De=n(l),ll=y(l,"P",{"data-svelte-h":!0}),i(ll)!=="svelte-lyz81m"&&(ll.innerHTML=Vs),Oe=n(l),p(el.$$.fragment,l),lM=n(l),Ml=y(l,"P",{"data-svelte-h":!0}),i(Ml)!=="svelte-163bvl5"&&(Ml.innerHTML=Bs),eM=n(l),sl=y(l,"UL",{"data-svelte-h":!0}),i(sl)!=="svelte-3el8k"&&(sl.innerHTML=ks),MM=n(l),tl=y(l,"P",{"data-svelte-h":!0}),i(tl)!=="svelte-10ra3tv"&&(tl.textContent=Gs),sM=n(l),p(nl.$$.fragment,l),tM=n(l),p(al.$$.fragment,l),nM=n(l),p(yl.$$.fragment,l),aM=n(l),p(il.$$.fragment,l),yM=n(l),p(Ul.$$.fragment,l),iM=n(l),p(pl.$$.fragment,l),UM=n(l),p(Tl.$$.fragment,l),pM=n(l),ol=y(l,"P",{"data-svelte-h":!0}),i(ol)!=="svelte-1r6d489"&&(ol.innerHTML=Xs),TM=n(l),p(Jl.$$.fragment,l),oM=n(l),Cl=y(l,"P",{"data-svelte-h":!0}),i(Cl)!=="svelte-hs3no7"&&(Cl.innerHTML=vs),JM=n(l),p(ml.$$.fragment,l),CM=n(l),wl=y(l,"P",{"data-svelte-h":!0}),i(wl)!=="svelte-12q3gxe"&&(wl.innerHTML=$s),mM=n(l),p(rl.$$.fragment,l),wM=n(l),p(jl.$$.fragment,l),rM=n(l),dl=y(l,"UL",{"data-svelte-h":!0}),i(dl)!=="svelte-n3fz09"&&(dl.innerHTML=Rs),jM=n(l),p(ul.$$.fragment,l),dM=n(l),cl=y(l,"UL",{"data-svelte-h":!0}),i(cl)!=="svelte-zppzms"&&(cl.innerHTML=Qs),uM=n(l),p(Il.$$.fragment,l),cM=n(l),fl=y(l,"UL",{"data-svelte-h":!0}),i(fl)!=="svelte-h3sif8"&&(fl.innerHTML=Ys),IM=n(l),p(Al.$$.fragment,l),fM=n(l),hl=y(l,"UL",{"data-svelte-h":!0}),i(hl)!=="svelte-18wpz97"&&(hl.innerHTML=Fs),AM=n(l),p(bl.$$.fragment,l),hM=n(l),Zl=y(l,"P",{"data-svelte-h":!0}),i(Zl)!=="svelte-1gtlfnj"&&(Zl.innerHTML=Ns),bM=n(l),p(_l.$$.fragment,l),ZM=n(l),Wl=y(l,"P",{"data-svelte-h":!0}),i(Wl)!=="svelte-oqrl65"&&(Wl.innerHTML=xs),_M=n(l),gl=y(l,"P",{"data-svelte-h":!0}),i(gl)!=="svelte-f0o4jc"&&(gl.innerHTML=zs),WM=n(l),p(Vl.$$.fragment,l),gM=n(l),Bl=y(l,"P",{"data-svelte-h":!0}),i(Bl)!=="svelte-cj2f57"&&(Bl.innerHTML=Es),VM=n(l),kl=y(l,"UL",{"data-svelte-h":!0}),i(kl)!=="svelte-is5pbj"&&(kl.innerHTML=Ss),BM=n(l),Gl=y(l,"P",{"data-svelte-h":!0}),i(Gl)!=="svelte-1p913uj"&&(Gl.textContent=Ls),kM=n(l),Xl=y(l,"UL",{"data-svelte-h":!0}),i(Xl)!=="svelte-13v9552"&&(Xl.innerHTML=Hs),GM=n(l),vl=y(l,"P",{"data-svelte-h":!0}),i(vl)!=="svelte-1psu9rx"&&(vl.textContent=qs),XM=n(l),p($l.$$.fragment,l),vM=n(l),Rl=y(l,"UL",{"data-svelte-h":!0}),i(Rl)!=="svelte-grvokw"&&(Rl.innerHTML=Ps),$M=n(l),p(Ql.$$.fragment,l),RM=n(l),Yl=y(l,"P",{"data-svelte-h":!0}),i(Yl)!=="svelte-1xgz07a"&&(Yl.innerHTML=Ks),QM=n(l),Fl=y(l,"P",{"data-svelte-h":!0}),i(Fl)!=="svelte-skk8sm"&&(Fl.innerHTML=Ds),YM=n(l),p(Nl.$$.fragment,l),FM=n(l),xl=y(l,"P",{"data-svelte-h":!0}),i(xl)!=="svelte-1c9b85z"&&(xl.textContent=Os),NM=n(l),zl=y(l,"UL",{"data-svelte-h":!0}),i(zl)!=="svelte-15o1gbp"&&(zl.innerHTML=lt),xM=n(l),El=y(l,"P",{"data-svelte-h":!0}),i(El)!=="svelte-1l6upaz"&&(El.textContent=et),zM=n(l),p(Sl.$$.fragment,l),EM=n(l),p(Ll.$$.fragment,l),SM=n(l),Hl=y(l,"P",{"data-svelte-h":!0}),i(Hl)!=="svelte-1p3hpdo"&&(Hl.innerHTML=Mt),LM=n(l),ql=y(l,"P",{"data-svelte-h":!0}),i(ql)!=="svelte-i621bo"&&(ql.innerHTML=st),HM=n(l),p(Pl.$$.fragment,l),qM=n(l),Kl=y(l,"P",{"data-svelte-h":!0}),i(Kl)!=="svelte-1c9b85z"&&(Kl.textContent=tt),PM=n(l),Dl=y(l,"UL",{"data-svelte-h":!0}),i(Dl)!=="svelte-1pogip"&&(Dl.innerHTML=nt),KM=n(l),p(Ol.$$.fragment,l),DM=n(l),le=y(l,"P",{"data-svelte-h":!0}),i(le)!=="svelte-6i8ef"&&(le.textContent=at),OM=n(l),p(ee.$$.fragment,l),ls=n(l),p(Me.$$.fragment,l),es=n(l),se=y(l,"P",{"data-svelte-h":!0}),i(se)!=="svelte-1lc4hka"&&(se.textContent=yt),Ms=n(l),te=y(l,"OL",{"data-svelte-h":!0}),i(te)!=="svelte-12pvc9t"&&(te.innerHTML=it),ss=n(l),ne=y(l,"P",{"data-svelte-h":!0}),i(ne)!=="svelte-ixzgc9"&&(ne.textContent=Ut),ts=n(l),p(ae.$$.fragment,l),ns=n(l),ye=y(l,"UL",{"data-svelte-h":!0}),i(ye)!=="svelte-1g0k1y8"&&(ye.innerHTML=pt),as=n(l),p(ie.$$.fragment,l),ys=n(l),Ue=y(l,"P",{"data-svelte-h":!0}),i(Ue)!=="svelte-1brs14g"&&(Ue.textContent=Tt),is=n(l),pe=y(l,"UL",{"data-svelte-h":!0}),i(pe)!=="svelte-1ep1va0"&&(pe.innerHTML=ot),Us=n(l),Te=y(l,"P",{"data-svelte-h":!0}),i(Te)!=="svelte-1q4t712"&&(Te.textContent=Jt),ps=n(l),oe=y(l,"UL",{"data-svelte-h":!0}),i(oe)!=="svelte-jpnk8d"&&(oe.innerHTML=Ct),Ts=n(l),Ce=y(l,"P",{}),mt(Ce).forEach(M),this.h()},h(){wt(r,"name","hf:doc:metadata"),wt(r,"content",bt)},m(l,e){ft(document.head,r),s(l,me,e),s(l,Je,e),s(l,we,e),T(j,l,e),s(l,re,e),T(d,l,e),s(l,je,e),s(l,u,e),s(l,de,e),T(c,l,e),s(l,ue,e),T(I,l,e),s(l,ce,e),s(l,f,e),s(l,Ie,e),s(l,A,e),s(l,fe,e),T(h,l,e),s(l,Ae,e),s(l,b,e),s(l,he,e),s(l,Z,e),s(l,be,e),s(l,_,e),s(l,Ze,e),s(l,W,e),s(l,_e,e),s(l,g,e),s(l,We,e),T(V,l,e),s(l,ge,e),s(l,B,e),s(l,Ve,e),s(l,k,e),s(l,Be,e),T(G,l,e),s(l,ke,e),T(X,l,e),s(l,Ge,e),s(l,v,e),s(l,Xe,e),s(l,$,e),s(l,ve,e),T(R,l,e),s(l,$e,e),T(Q,l,e),s(l,Re,e),s(l,Y,e),s(l,Qe,e),T(F,l,e),s(l,Ye,e),T(N,l,e),s(l,Fe,e),T(x,l,e),s(l,Ne,e),T(z,l,e),s(l,xe,e),T(E,l,e),s(l,ze,e),s(l,S,e),s(l,Ee,e),T(L,l,e),s(l,Se,e),T(H,l,e),s(l,Le,e),s(l,q,e),s(l,He,e),s(l,P,e),s(l,qe,e),T(K,l,e),s(l,Pe,e),s(l,D,e),s(l,Ke,e),T(O,l,e),s(l,De,e),s(l,ll,e),s(l,Oe,e),T(el,l,e),s(l,lM,e),s(l,Ml,e),s(l,eM,e),s(l,sl,e),s(l,MM,e),s(l,tl,e),s(l,sM,e),T(nl,l,e),s(l,tM,e),T(al,l,e),s(l,nM,e),T(yl,l,e),s(l,aM,e),T(il,l,e),s(l,yM,e),T(Ul,l,e),s(l,iM,e),T(pl,l,e),s(l,UM,e),T(Tl,l,e),s(l,pM,e),s(l,ol,e),s(l,TM,e),T(Jl,l,e),s(l,oM,e),s(l,Cl,e),s(l,JM,e),T(ml,l,e),s(l,CM,e),s(l,wl,e),s(l,mM,e),T(rl,l,e),s(l,wM,e),T(jl,l,e),s(l,rM,e),s(l,dl,e),s(l,jM,e),T(ul,l,e),s(l,dM,e),s(l,cl,e),s(l,uM,e),T(Il,l,e),s(l,cM,e),s(l,fl,e),s(l,IM,e),T(Al,l,e),s(l,fM,e),s(l,hl,e),s(l,AM,e),T(bl,l,e),s(l,hM,e),s(l,Zl,e),s(l,bM,e),T(_l,l,e),s(l,ZM,e),s(l,Wl,e),s(l,_M,e),s(l,gl,e),s(l,WM,e),T(Vl,l,e),s(l,gM,e),s(l,Bl,e),s(l,VM,e),s(l,kl,e),s(l,BM,e),s(l,Gl,e),s(l,kM,e),s(l,Xl,e),s(l,GM,e),s(l,vl,e),s(l,XM,e),T($l,l,e),s(l,vM,e),s(l,Rl,e),s(l,$M,e),T(Ql,l,e),s(l,RM,e),s(l,Yl,e),s(l,QM,e),s(l,Fl,e),s(l,YM,e),T(Nl,l,e),s(l,FM,e),s(l,xl,e),s(l,NM,e),s(l,zl,e),s(l,xM,e),s(l,El,e),s(l,zM,e),T(Sl,l,e),s(l,EM,e),T(Ll,l,e),s(l,SM,e),s(l,Hl,e),s(l,LM,e),s(l,ql,e),s(l,HM,e),T(Pl,l,e),s(l,qM,e),s(l,Kl,e),s(l,PM,e),s(l,Dl,e),s(l,KM,e),T(Ol,l,e),s(l,DM,e),s(l,le,e),s(l,OM,e),T(ee,l,e),s(l,ls,e),T(Me,l,e),s(l,es,e),s(l,se,e),s(l,Ms,e),s(l,te,e),s(l,ss,e),s(l,ne,e),s(l,ts,e),T(ae,l,e),s(l,ns,e),s(l,ye,e),s(l,as,e),T(ie,l,e),s(l,ys,e),s(l,Ue,e),s(l,is,e),s(l,pe,e),s(l,Us,e),s(l,Te,e),s(l,ps,e),s(l,oe,e),s(l,Ts,e),s(l,Ce,e),os=!0},p:jt,i(l){os||(o(j.$$.fragment,l),o(d.$$.fragment,l),o(c.$$.fragment,l),o(I.$$.fragment,l),o(h.$$.fragment,l),o(V.$$.fragment,l),o(G.$$.fragment,l),o(X.$$.fragment,l),o(R.$$.fragment,l),o(Q.$$.fragment,l),o(F.$$.fragment,l),o(N.$$.fragment,l),o(x.$$.fragment,l),o(z.$$.fragment,l),o(E.$$.fragment,l),o(L.$$.fragment,l),o(H.$$.fragment,l),o(K.$$.fragment,l),o(O.$$.fragment,l),o(el.$$.fragment,l),o(nl.$$.fragment,l),o(al.$$.fragment,l),o(yl.$$.fragment,l),o(il.$$.fragment,l),o(Ul.$$.fragment,l),o(pl.$$.fragment,l),o(Tl.$$.fragment,l),o(Jl.$$.fragment,l),o(ml.$$.fragment,l),o(rl.$$.fragment,l),o(jl.$$.fragment,l),o(ul.$$.fragment,l),o(Il.$$.fragment,l),o(Al.$$.fragment,l),o(bl.$$.fragment,l),o(_l.$$.fragment,l),o(Vl.$$.fragment,l),o($l.$$.fragment,l),o(Ql.$$.fragment,l),o(Nl.$$.fragment,l),o(Sl.$$.fragment,l),o(Ll.$$.fragment,l),o(Pl.$$.fragment,l),o(Ol.$$.fragment,l),o(ee.$$.fragment,l),o(Me.$$.fragment,l),o(ae.$$.fragment,l),o(ie.$$.fragment,l),os=!0)},o(l){J(j.$$.fragment,l),J(d.$$.fragment,l),J(c.$$.fragment,l),J(I.$$.fragment,l),J(h.$$.fragment,l),J(V.$$.fragment,l),J(G.$$.fragment,l),J(X.$$.fragment,l),J(R.$$.fragment,l),J(Q.$$.fragment,l),J(F.$$.fragment,l),J(N.$$.fragment,l),J(x.$$.fragment,l),J(z.$$.fragment,l),J(E.$$.fragment,l),J(L.$$.fragment,l),J(H.$$.fragment,l),J(K.$$.fragment,l),J(O.$$.fragment,l),J(el.$$.fragment,l),J(nl.$$.fragment,l),J(al.$$.fragment,l),J(yl.$$.fragment,l),J(il.$$.fragment,l),J(Ul.$$.fragment,l),J(pl.$$.fragment,l),J(Tl.$$.fragment,l),J(Jl.$$.fragment,l),J(ml.$$.fragment,l),J(rl.$$.fragment,l),J(jl.$$.fragment,l),J(ul.$$.fragment,l),J(Il.$$.fragment,l),J(Al.$$.fragment,l),J(bl.$$.fragment,l),J(_l.$$.fragment,l),J(Vl.$$.fragment,l),J($l.$$.fragment,l),J(Ql.$$.fragment,l),J(Nl.$$.fragment,l),J(Sl.$$.fragment,l),J(Ll.$$.fragment,l),J(Pl.$$.fragment,l),J(Ol.$$.fragment,l),J(ee.$$.fragment,l),J(Me.$$.fragment,l),J(ae.$$.fragment,l),J(ie.$$.fragment,l),os=!1},d(l){l&&(M(me),M(Je),M(we),M(re),M(je),M(u),M(de),M(ue),M(ce),M(f),M(Ie),M(A),M(fe),M(Ae),M(b),M(he),M(Z),M(be),M(_),M(Ze),M(W),M(_e),M(g),M(We),M(ge),M(B),M(Ve),M(k),M(Be),M(ke),M(Ge),M(v),M(Xe),M($),M(ve),M($e),M(Re),M(Y),M(Qe),M(Ye),M(Fe),M(Ne),M(xe),M(ze),M(S),M(Ee),M(Se),M(Le),M(q),M(He),M(P),M(qe),M(Pe),M(D),M(Ke),M(De),M(ll),M(Oe),M(lM),M(Ml),M(eM),M(sl),M(MM),M(tl),M(sM),M(tM),M(nM),M(aM),M(yM),M(iM),M(UM),M(pM),M(ol),M(TM),M(oM),M(Cl),M(JM),M(CM),M(wl),M(mM),M(wM),M(rM),M(dl),M(jM),M(dM),M(cl),M(uM),M(cM),M(fl),M(IM),M(fM),M(hl),M(AM),M(hM),M(Zl),M(bM),M(ZM),M(Wl),M(_M),M(gl),M(WM),M(gM),M(Bl),M(VM),M(kl),M(BM),M(Gl),M(kM),M(Xl),M(GM),M(vl),M(XM),M(vM),M(Rl),M($M),M(RM),M(Yl),M(QM),M(Fl),M(YM),M(FM),M(xl),M(NM),M(zl),M(xM),M(El),M(zM),M(EM),M(SM),M(Hl),M(LM),M(ql),M(HM),M(qM),M(Kl),M(PM),M(Dl),M(KM),M(DM),M(le),M(OM),M(ls),M(es),M(se),M(Ms),M(te),M(ss),M(ne),M(ts),M(ns),M(ye),M(as),M(ys),M(Ue),M(is),M(pe),M(Us),M(Te),M(ps),M(oe),M(Ts),M(Ce)),M(r),C(j,l),C(d,l),C(c,l),C(I,l),C(h,l),C(V,l),C(G,l),C(X,l),C(R,l),C(Q,l),C(F,l),C(N,l),C(x,l),C(z,l),C(E,l),C(L,l),C(H,l),C(K,l),C(O,l),C(el,l),C(nl,l),C(al,l),C(yl,l),C(il,l),C(Ul,l),C(pl,l),C(Tl,l),C(Jl,l),C(ml,l),C(rl,l),C(jl,l),C(ul,l),C(Il,l),C(Al,l),C(bl,l),C(_l,l),C(Vl,l),C($l,l),C(Ql,l),C(Nl,l),C(Sl,l),C(Ll,l),C(Pl,l),C(Ol,l),C(ee,l),C(Me,l),C(ae,l),C(ie,l)}}}const bt='{"title":"Contributing Custom Models for Training","local":"contributing-custom-models-for-training","sections":[{"title":"Architecture Components","local":"architecture-components","sections":[{"title":"1. NeuronModelMixin","local":"1-neuronmodelmixin","sections":[],"depth":3},{"title":"2. Weight Transformation Specs","local":"2-weight-transformation-specs","sections":[],"depth":3},{"title":"3. Parallel Layers","local":"3-parallel-layers","sections":[],"depth":3}],"depth":2},{"title":"Implementation Steps","local":"implementation-steps","sections":[{"title":"Step 1: Create Model Structure","local":"step-1-create-model-structure","sections":[],"depth":3},{"title":"Step 2: Implement the Model Building Blocks","local":"step-2-implement-the-model-building-blocks","sections":[{"title":"Imports and Dependencies","local":"imports-and-dependencies","sections":[],"depth":4},{"title":"Embedding Layer","local":"embedding-layer","sections":[],"depth":4},{"title":"MLP Layer with Fused Linears","local":"mlp-layer-with-fused-linears","sections":[],"depth":4},{"title":"Attention Layer","local":"attention-layer","sections":[],"depth":4}],"depth":3},{"title":"Step 3: Implement Main Model Classes","local":"step-3-implement-main-model-classes","sections":[{"title":"Base Model","local":"base-model","sections":[],"depth":4},{"title":"CausalLM Model","local":"causallm-model","sections":[],"depth":4}],"depth":3},{"title":"Step 4: Register Model","local":"step-4-register-model","sections":[],"depth":3}],"depth":2},{"title":"Best Practices","local":"best-practices","sections":[{"title":"1. Parallel Layer Configuration","local":"1-parallel-layer-configuration","sections":[],"depth":3},{"title":"2. Weight Transformation Specs","local":"2-weight-transformation-specs","sections":[],"depth":3},{"title":"3. Pipeline Parallelism","local":"3-pipeline-parallelism","sections":[],"depth":3},{"title":"4. Flash Attention Support","local":"4-flash-attention-support","sections":[],"depth":3}],"depth":2},{"title":"Testing Your Implementation","local":"testing-your-implementation","sections":[{"title":"1. Custom Modeling Validation","local":"1-custom-modeling-validation","sections":[],"depth":3},{"title":"2. End-to-End Training Validation","local":"2-end-to-end-training-validation","sections":[],"depth":3},{"title":"3. Auto Model Loading","local":"3-auto-model-loading","sections":[],"depth":3},{"title":"4. Running Tests","local":"4-running-tests","sections":[],"depth":3},{"title":"5. Test Requirements","local":"5-test-requirements","sections":[],"depth":3}],"depth":2},{"title":"Common Issues","local":"common-issues","sections":[],"depth":2},{"title":"Additional Resources","local":"additional-resources","sections":[],"depth":2}],"depth":1}';function Zt(Js){return dt(()=>{new URLSearchParams(window.location.search).get("fw")}),[]}class kt extends ut{constructor(r){super(),ct(this,r,Zt,ht,rt,{})}}export{kt as component};