File tree Expand file tree Collapse file tree
nemo_automodel/recipes/llm Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -163,36 +163,6 @@ def _get_num_thd_chunks(pp_enabled, cfg):
163163 return 1
164164
165165
166- def _mtp_is_enabled (cfg , model_parts ) -> bool :
167- """Return True if Multi-Token Prediction is enabled for this run.
168-
169- Checks both signals because either may be missing depending on how the
170- model was constructed:
171-
172- * YAML override / explicit DeepseekV4Config: the
173- ``model.config.num_nextn_predict_layers`` field is the user-facing
174- knob and is present on the cfg before any model is built.
175- * Constructed model: V4's ``ForCausalLM.__init__`` materializes
176- ``self.mtp_config``. Walking ``modules()`` catches it on the root
177- or on any submodule that retained the attribute after wrapping.
178-
179- The module walk alone isn't sufficient: pipeline-parallel wrapping can
180- replace the V4 root with a stage container that no longer exposes
181- ``mtp_config``, in which case only the cfg lookup catches MTP.
182- """
183- n = int (cfg .get ("model.config.num_nextn_predict_layers" , 0 ) or 0 )
184- if n > 0 :
185- return True
186- for mp in model_parts :
187- if mp is None :
188- continue
189- for sub in mp .modules ():
190- mc = getattr (sub , "mtp_config" , None )
191- if mc is not None and getattr (mc , "enabled" , False ):
192- return True
193- return False
194-
195-
196166def build_model (
197167 cfg_model ,
198168 cfg_peft ,
You can’t perform that action at this time.
0 commit comments