refactor: sanitize tag usage

Author: dacorvo

optimum/neuron/models/inference/backend/graph_builder.py

@@ -19,9 +19,8 @@
 
 
 class NxDGraphBuilder(ABC):
-    def __init__(self, tag: str, priority_model_idx: int):
+    def __init__(self, priority_model_idx: int):
         super().__init__()
-        self.tag = tag
         self.priority_model_idx = priority_model_idx
 
     @abstractmethod

optimum/neuron/models/inference/backend/modules/decoder/decoder_builder.py

@@ -30,10 +30,9 @@ def __init__(
         max_tokens: int,
         active_tokens: int,
         model_cls,
-        tag="",
         priority_model_idx: int = None,
     ) -> None:
-        super().__init__(tag, priority_model_idx)
+        super().__init__(priority_model_idx)
         self.config = config
         self.neuron_config = neuron_config
         self.max_tokens = max_tokens

optimum/neuron/models/inference/backend/modules/decoder/modeling_decoder.py

@@ -384,69 +384,36 @@ def _create_speculation_config(neuron_config: NxDNeuronConfig) -> NxDNeuronConfi
         return spec_neuron_config
 
     @staticmethod
-    def _create_context_encoding_builder(model_cls, config, neuron_config):
+    def create_graph_builders(model_cls, config, neuron_config):
+        graph_builders = {}
         ctx_neuron_config = NxDModelForCausalLM._create_context_encoding_config(neuron_config)
-
-        return NxDDecoderBuilder(
+        graph_builders["context_encoding"] = NxDDecoderBuilder(
             config=config,
             neuron_config=ctx_neuron_config,
             max_tokens=ctx_neuron_config.max_context_length,
             active_tokens=ctx_neuron_config.max_context_length,
             model_cls=model_cls,
-            tag=CONTEXT_ENCODING_MODEL_TAG,
         )
-
-    @staticmethod
-    def _create_token_generation_builder(model_cls, config, neuron_config, enable_wlt_optimization: bool = True):
         tkg_neuron_config = NxDModelForCausalLM._create_token_generation_config(neuron_config)
-
-        return NxDDecoderBuilder(
+        graph_builders["token_generation"] = NxDDecoderBuilder(
             config=config,
             neuron_config=tkg_neuron_config,
             max_tokens=tkg_neuron_config.sequence_length,
             active_tokens=1,
             model_cls=model_cls,
-            tag=TOKEN_GENERATION_MODEL_TAG,
-            priority_model_idx=0 if enable_wlt_optimization else None,  # to turn on weight layout optimization
-        )
-
-    @staticmethod
-    def _create_speculation_builder(model_cls, config, neuron_config):
-        spec_neuron_config = NxDModelForCausalLM._create_speculation_config(neuron_config)
-
-        return NxDDecoderBuilder(
-            config=config,
-            neuron_config=spec_neuron_config,
-            max_tokens=spec_neuron_config.sequence_length,
-            active_tokens=spec_neuron_config.speculation_length,
-            model_cls=model_cls,
-            tag=SPECULATION_MODEL_TAG,
             priority_model_idx=0,  # to turn on weight layout optimization
         )
-
-    @staticmethod
-    def create_model_builders(model_cls, config, neuron_config):
-        model_builders = [
-            NxDModelForCausalLM._create_context_encoding_builder(
-                model_cls,
-                config,
-                neuron_config,
-            ),
-            NxDModelForCausalLM._create_token_generation_builder(
-                model_cls,
-                config,
-                neuron_config,
-            ),
-        ]
         if neuron_config.speculation_length > 0:
-            model_builders.append(
-                NxDModelForCausalLM._create_speculation_builder(
-                    model_cls,
-                    config,
-                    neuron_config,
-                )
+            spec_neuron_config = NxDModelForCausalLM._create_speculation_config(neuron_config)
+            graph_builders["speculation_model"] = NxDDecoderBuilder(
+                config=config,
+                neuron_config=spec_neuron_config,
+                max_tokens=spec_neuron_config.sequence_length,
+                active_tokens=spec_neuron_config.speculation_length,
+                model_cls=model_cls,
+                priority_model_idx=0,  # to turn on weight layout optimization
             )
-        return model_builders
+        return graph_builders
 
     def forward(
         self,
@@ -650,7 +617,7 @@ def _from_pretrained(
                 traced_model = torch.jit.load(os.path.join(tmpdir, cls.COMPILED_MODEL_FILE_NAME))
         else:
             traced_model = torch.jit.load(os.path.join(model_id, cls.COMPILED_MODEL_FILE_NAME))
-        model_builders = NxDModelForCausalLM.create_model_builders(
+        model_builders = NxDModelForCausalLM.create_graph_builders(
             cls._model_cls, config=config, neuron_config=neuron_config
         )
         model = cls(
@@ -707,7 +674,7 @@ def _export(
         # Evaluate head_dim if it is defined but set to null (like in Mixtral for transformers 4.54+)
         if hasattr(config, "head_dim") and config.head_dim is None:
             config.head_dim = config.hidden_size // config.num_attention_heads
-        model_builders = cls.create_model_builders(
+        model_builders = cls.create_graph_builders(
             model_cls=cls._model_cls,
             config=config,
             neuron_config=neuron_config,

optimum/neuron/models/inference/backend/modules/kvcache/kv_cache_manager.py

@@ -93,7 +93,7 @@ def get_cache(self, seq_len: int, skip_slice=False, **kwargs):
         """
         Return network (all layers)'s previously cached K and V, up to seq_len.
 
-        :param seq_len: sequence length (or bucket size from auto-bucketing e.g. 128, 512, 1024 etc.)
+        :param seq_len: sequence length
         :return: list of tuple of (K, V)
         """
         slice_index, gather_index = None, None
@@ -129,9 +129,9 @@ def update_cache(
         :param scatter_index: tensor representing index to update
         :param is_for_context_encoding: bool
         :param seq_ids: tensor of size (batch_sz)
-        :param position_ids: tensor of size (batch_sz, bucket_sz)
+        :param position_ids: tensor of size (batch_sz, seq_len)
         :param new_key_values: list of tuple, the latest kv obtained at the end of the network from forward pass
-        :param seq_len: sequence length (or bucket size from auto-bucketing e.g. 128, 512, 1024 etc.)
+        :param seq_len: sequence length
         :param scatter_index: tensor representing index to update
         :param active_mask: tensor representing index to update
         :param kvcache_buffer: if passed key states are updates to this buffer.

optimum/neuron/models/inference/backend/pretrained_model.py

@@ -49,7 +49,7 @@ def get_shards_path(dest_path):
 
 def get_builder(
     neuron_config: NxDNeuronConfig,
-    model_wrappers: list[NxDGraphBuilder],
+    model_wrappers: dict[str, NxDGraphBuilder],
     debug: bool = False,
     checkpoint_loader=None,
     compiler_args: str = None,
@@ -86,9 +86,9 @@ def get_builder(
         logical_nc_config=neuron_config.logical_nc_config,
         weights_to_skip_layout_optimization=neuron_config.weights_to_skip_layout_optimization,
     )
-    for model in model_wrappers:
+    for tag, model in model_wrappers.items():
         builder.add(
-            key=model.tag,
+            key=tag,
             model_instance=model.get_model_instance(),
             example_inputs=model.input_generator(),
             compiler_args=compiler_args,
@@ -109,7 +109,7 @@ def __init__(
         config: PretrainedConfig,
         neuron_config: NxDNeuronConfig,
         traced_model: torch.jit.ScriptModule,
-        model_wrappers: list[NxDGraphBuilder],
+        model_wrappers: dict[str, NxDGraphBuilder],
     ):
         self.config = copy.deepcopy(config)
         self.neuron_config = copy.deepcopy(neuron_config)
@@ -131,7 +131,7 @@ def get_compiler_args(cls, neuron_config) -> str | None:
         return None
 
     @staticmethod
-    def compile(neuron_config, model_wrappers: list[NxDGraphBuilder], compiler_args: str, debug: bool = False):
+    def compile(neuron_config, model_wrappers: dict[str, NxDGraphBuilder], compiler_args: str, debug: bool = False):
         builder = get_builder(neuron_config, model_wrappers, debug=debug, compiler_args=compiler_args)
         return builder.trace(initialize_model_weights=False)