add preprocessing mlp for hstu (#98)

* add item mpl and contextual mlp

* fix

* fix ci

* fix

Author: shijieliu

.gitignore

@@ -29,4 +29,5 @@ embedding_checkpoint/
 .idea/
 .DS_Store
 *.pickle
-*.xlsx
+*.xlsx
+pcie_lookup_poc/*

examples/commons/utils/initialize.py

@@ -16,7 +16,13 @@
 import os
 
 import torch
-from megatron.core import parallel_state, tensor_parallel
+
+try:
+    from megatron.core import parallel_state, tensor_parallel
+except ImportError:
+    print("megatron.core is not installed, training is not supported.")
+    parallel_state = None
+    tensor_parallel = None
 
 
 def initialize_single_rank():

examples/hstu/benchmark_ranking.gin

@@ -38,6 +38,8 @@ BenchmarkDatasetArgs.contextual_feature_names=[]
 BenchmarkDatasetArgs.action_feature_name='action'
 BenchmarkDatasetArgs.max_num_candidates=0
 
+NetworkArgs.item_embedding_dim = 128
+NetworkArgs.contextual_embedding_dim = 256
 NetworkArgs.num_layers = 8
 NetworkArgs.num_attention_heads = 4
 NetworkArgs.hidden_size = 1024

examples/hstu/configs/__init__.py

@@ -2,6 +2,7 @@
 from .hstu_config import (
     HSTUConfig,
     HSTULayerType,
+    HSTUPreprocessingConfig,
     KernelBackend,
     PositionEncodingConfig,
     get_hstu_config,
@@ -29,6 +30,7 @@
     "task_config",
     "ConfigType",
     "PositionEncodingConfig",
+    "HSTUPreprocessingConfig",
     "HSTUConfig",
     "get_hstu_config",
     "RankingConfig",

examples/hstu/configs/hstu_config.py

@@ -77,12 +77,19 @@ class PositionEncodingConfig:
     use_time_encoding: bool
 
 
+@dataclass
+class HSTUPreprocessingConfig:
+    item_embedding_dim: int
+    contextual_embedding_dim: int
+
+
 @dataclass
 class HSTUConfig(TransformerConfig):
     """
     HSTUConfig is a configuration data class for the HSTU model, inheriting from TransformerConfig.
 
     Args:
+      hstu_preprocessing_config (HSTUPreprocessingConfig): HSTU preprocessing config. Defaults to None.
       position_encoding_config (PositionEncodingConfig): Position embedding config. Defaults to None.
       is_causal (bool): Indicates if the model is causal. Defaults to True.
       enable_relative_attention_bias (bool): Flag to enable relative attention bias. Defaults to False.
@@ -97,6 +104,7 @@ class HSTUConfig(TransformerConfig):
       recompute_input_silu (bool): Flag to enable recompute input silu. Defaults to False.
     """
 
+    hstu_preprocessing_config: Optional[HSTUPreprocessingConfig] = None
     position_encoding_config: Optional[PositionEncodingConfig] = None
     is_causal: bool = True
     enable_relative_attention_bias: bool = False
@@ -131,6 +139,7 @@ def get_hstu_config(
     num_attention_heads,
     num_layers,
     dtype,
+    hstu_preprocessing_config: Optional[HSTUPreprocessingConfig] = None,
     position_encoding_config: Optional[PositionEncodingConfig] = None,
     hidden_dropout=0.2,
     norm_epsilon=1e-5,
@@ -156,6 +165,7 @@ def get_hstu_config(
         num_attention_heads (int): Number of attention heads.
         num_layers (int): Number of attention layers.
         dtype (torch.dtype): Data type (e.g., torch.float16).
+        hstu_preprocessing_config (Optional[HSTUPreprocessingConfig], optional): HSTU preprocessing config. Defaults to None.
         position_encoding_config (Optional[PositionEncodingConfig], optional): Position embedding config. Defaults to None.
         hidden_dropout (float, optional): Dropout rate for hidden layers. Defaults to 0.2.
         norm_epsilon (float, optional): Epsilon value for normalization. Defaults to 1e-5.
@@ -181,6 +191,7 @@ def get_hstu_config(
         async_wgrad_stream = None
         async_wgrad_event = None
     return HSTUConfig(  # type: ignore
+        hstu_preprocessing_config=hstu_preprocessing_config,
         position_encoding_config=position_encoding_config,
         hidden_size=hidden_size,
         kv_channels=kv_channels,

examples/hstu/configs/inference_config.py

@@ -17,7 +17,7 @@
 
 import torch
 
-from .hstu_config import PositionEncodingConfig
+from .hstu_config import HSTUPreprocessingConfig, PositionEncodingConfig
 
 
 @dataclass
@@ -156,6 +156,7 @@ class InferenceHSTUConfig:
     is_causal: bool = True
     target_group_size: int = 1
     position_encoding_config: Optional[PositionEncodingConfig] = None
+    hstu_preprocessing_config: Optional[HSTUPreprocessingConfig] = None
 
     def __post_init__(self):
         assert self.is_causal

examples/hstu/dataset/utils.py

@@ -77,6 +77,32 @@ def __post_init__(self):
         )
         assert isinstance(self.max_num_candidates, int)
 
+    def to(self, device: torch.device, non_blocking: bool = False) -> "Batch":  # type: ignore
+        """
+        Move the batch to the specified device.
+
+        Args:
+            device (torch.device): The device to move the batch to.
+            non_blocking (bool, optional): Whether to perform the move asynchronously. Defaults to False.
+
+        Returns:
+            RankingBatch: The batch on the specified device.
+        """
+        return Batch(
+            features=self.features.to(device=device, non_blocking=non_blocking),
+            batch_size=self.batch_size,
+            feature_to_max_seqlen=self.feature_to_max_seqlen,
+            contextual_feature_names=self.contextual_feature_names,
+            item_feature_name=self.item_feature_name,
+            action_feature_name=self.action_feature_name,
+            max_num_candidates=self.max_num_candidates,
+            num_candidates=self.num_candidates.to(
+                device=device, non_blocking=non_blocking
+            )
+            if self.num_candidates is not None
+            else None,
+        )
+
     @staticmethod
     def random(
         batch_size: int,

examples/hstu/model/inference_ranking_gr.py

@@ -358,7 +358,7 @@ def forward(
     ):
         with torch.inference_mode():
             kvcache_metadata = self.prepare_kv_cache(batch, user_ids, user_start_pos)
-            jagged_data = self._hstu_block.hstu_preprocess(
+            jagged_data = self._hstu_block._preprocessor(
                 embeddings=self._embedding_collection(batch.features),
                 batch=batch,
             )
@@ -400,7 +400,7 @@ def forward(
                 torch.cuda.current_stream()
             )
 
-            jagged_data = self._hstu_block.hstu_postprocess(jagged_data)
+            jagged_data = self._hstu_block._postprocessor(jagged_data)
             jagged_item_logit = self._dense_module(jagged_data.values)
             self._offload_states = self.offload_kv_cache_async(
                 user_ids, kvcache_metadata

examples/hstu/model/ranking_gr.py

@@ -50,17 +50,11 @@ def __init__(
         self._hstu_config = hstu_config
         self._task_config = task_config
 
-        self._embedding_dim = hstu_config.hidden_size
-        for ebc_config in task_config.embedding_configs:
-            assert (
-                ebc_config.dim == self._embedding_dim
-            ), "hstu layer hidden size should equal to embedding dim"
-
         self._embedding_collection = ShardedEmbedding(task_config.embedding_configs)
 
         self._hstu_block = HSTUBlock(hstu_config)
         self._mlp = MLP(
-            self._embedding_dim,
+            hstu_config.hidden_size,
             task_config.prediction_head_arch,
             task_config.prediction_head_act_type,
             task_config.prediction_head_bias,

examples/hstu/modules/hstu_block.py

@@ -1,6 +1,6 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 
-from typing import Dict, Optional, Union
+from typing import Dict, Union
 
 import torch
 from commons.utils.nvtx_op import output_nvtx_hook
@@ -9,14 +9,9 @@
 from megatron.core.transformer.module import MegatronModule
 from modules.debug.debug_hstu_layer import HSTULayer as DebugHSTULayer
 from modules.fused_hstu_layer import FusedHSTULayer
+from modules.hstu_processor import HSTUBlockPostprocessor, HSTUBlockPreprocessor
 from modules.jagged_data import JaggedData
 from modules.native_hstu_layer import HSTULayer as NativeHSTULayer
-from modules.position_encoder import HSTUPositionalEncoder
-from modules.utils import hstu_postprocess_embeddings, hstu_preprocess_embeddings
-from ops.triton_ops.triton_jagged import (  # type: ignore[attr-defined]
-    triton_concat_2D_jagged,
-    triton_split_2D_jagged,
-)
 from torchrec.sparse.jagged_tensor import JaggedTensor
 
 
@@ -39,16 +34,9 @@ def __init__(
         if self.config.fp16:
             self._training_dtype = torch.float16
 
-        self._positional_encoder: Optional[HSTUPositionalEncoder] = None
-        if config.position_encoding_config is not None:
-            self._positional_encoder = HSTUPositionalEncoder(
-                num_position_buckets=config.position_encoding_config.num_position_buckets,
-                num_time_buckets=config.position_encoding_config.num_time_buckets,
-                embedding_dim=config.hidden_size,
-                is_inference=False,
-                use_time_encoding=config.position_encoding_config.use_time_encoding,
-                training_dtype=self._training_dtype,
-            )
+        self._preprocessor = HSTUBlockPreprocessor(config, is_inference=False)
+        self._postprocessor = HSTUBlockPostprocessor(is_inference=False)
+
         HSTULayerImpl = (
             FusedHSTULayer
             if config.hstu_layer_type == HSTULayerType.FUSED
@@ -59,62 +47,6 @@ def __init__(
         self._attention_layers = torch.nn.ModuleList(
             [HSTULayerImpl(config) for l in range(self.config.num_layers)]
         )
-        self._dropout_ratio = config.hidden_dropout
-
-    @output_nvtx_hook(nvtx_tag="HSTUBlock preprocess", hook_key_or_attr_name="values")
-    def hstu_preprocess(
-        self, embeddings: Dict[str, JaggedTensor], batch: RankingBatch
-    ) -> JaggedData:
-        """
-        Preprocesses the embeddings for use in the HSTU architecture.
-
-        This method performs the following steps:
-        1. **Interleaving**: If action embeddings are present, interleaves them with item embeddings.
-        2. **Concatenation**: Concatenates contextual, item, and action embeddings for each sample, following the order specified in the batch.
-        3. **Position Encoding**: Applies position encoding to the concatenated embeddings.
-
-        Args:
-            embeddings (Dict[str, JaggedTensor]): A dictionary of embeddings where each key corresponds to a feature name and the value is a jagged tensor.
-            batch (RankingBatch): The batch of ranking data.
-
-        Returns:
-            JaggedData: The preprocessed jagged data, ready for further processing in the HSTU architecture.
-        """
-        # Interleaving & concatenation
-        jd = hstu_preprocess_embeddings(embeddings, batch, is_inference=False)
-
-        if self._positional_encoder is not None:
-            jd.values = self._positional_encoder(
-                max_seq_len=jd.max_seqlen,
-                seq_lengths=jd.seqlen,
-                seq_offsets=jd.seqlen_offsets,
-                seq_timestamps=None,
-                seq_embeddings=jd.values,
-                num_targets=jd.num_candidates,
-            )
-
-        jd.values = torch.nn.functional.dropout(
-            jd.values,
-            p=self._dropout_ratio,
-            training=self.training,
-        ).to(self._training_dtype)
-        return jd
-
-    @output_nvtx_hook(nvtx_tag="HSTUBlock postprocess", hook_key_or_attr_name="values")
-    def hstu_postprocess(self, jd: JaggedData) -> JaggedData:
-        """
-        Postprocess the output from the HSTU architecture.
-        1. If max_num_candidates > 0, split and only keep last ``num_candidates`` embeddings as candidates embedding for further processing.
-        2. Remove action embeddings if present. Only use item embedding for further processing.
-
-        Args:
-            jd (JaggedData): The jagged data output from the HSTU architecture that needs further processing.
-
-        Returns:
-            JaggedData: The postprocessed jagged data.
-        """
-
-        return hstu_postprocess_embeddings(jd, is_inference=False)
 
     @output_nvtx_hook(nvtx_tag="HSTUBlock", hook_key_or_attr_name="values")
     def forward(
@@ -132,7 +64,7 @@ def forward(
         Returns:
             JaggedData: The output jagged data.
         """
-        jd = self.hstu_preprocess(embeddings, batch)
+        jd = self._preprocessor(embeddings, batch)
         for hstu_layer in self._attention_layers:
             jd = hstu_layer(jd)
-        return self.hstu_postprocess(jd)
+        return self._postprocessor(jd)