[perf] Improve performance for putting jagged tensor (Ascend#36)

## Background

When users input a TensorDict containing jagged tensors (nested
tensors), the `put_data` process becomes extremely slow.

Specifically, the `_filter_storage_data` function uses
`itemgetter(*batch_indexes)(data[fname])` to extract individual items
from each tensor in the TensorDict. This indexing approach works
efficiently for strided tensors but is extremely inefficient for jagged
tensors.

## Root Cause
For jagged tensors, itemgetter with multiple batch indexes requires
repeated indexing operations, which is $\mathcal{O}(n)$ for each access.
When extracting multiple samples, this becomes $\mathcal{O}(n²)$
complexity.


## Solution

We unbind nested tensor before accessing each sample from it.

```python3
  # unbind nested tensor
  results: dict = {}
  for field in sorted(data.keys()):
      field_data = data[field]
      if isinstance(field_data, Tensor) and field_data.is_nested:
          results[field] = field_data.unbind()
      else:
          results[field] = field_data
```



## Simple Reproduction Script

```python3
  import torch
  import time
  from operator import itemgetter

  # Create a jagged tensor with 1000 samples
  offsets = torch.tensor([0] + list(torch.randint(10, 50, (1001,)).cumsum(0)))
  values = torch.randn(offsets[-1].item(), 128)
  jagged = torch.nested.as_nested_tensor(
      [values[offsets[i]:offsets[i+1]] for i in range(1000)],
      layout=torch.jagged
  )

  batch_indexes = list(range(0, 1000, 10))  # 100 indexes

  # Method 1: Direct itemgetter on jagged tensor (SLOW)
  start = time.perf_counter()
  result = itemgetter(*batch_indexes)(jagged)
  print(f"Direct itemgetter: {(time.perf_counter() - start)*1000:.2f} ms")

  # Method 2: Unbind first, then itemgetter (FAST)
  start = time.perf_counter()
  field_list = jagged.unbind()
  result = itemgetter(*batch_indexes)(field_list)
  print(f"Unbind + itemgetter: {(time.perf_counter() - start)*1000:.2f} ms")

```
Output: 
```bash
Direct itemgetter: 150.94 ms
Unbind + itemgetter: 1.80 ms
```

---------

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Author: 0oshowero0

transfer_queue/metadata.py

@@ -26,6 +26,7 @@
 import torch
 from tensordict import TensorDict
 from tensordict.tensorclass import NonTensorData, NonTensorStack
+from torch import Tensor
 
 from transfer_queue.utils.enum_utils import ProductionStatus
 
@@ -815,18 +816,26 @@ def _extract_field_metas(tensor_dict: TensorDict, set_all_ready: bool = True) ->
 
     production_status = ProductionStatus.READY_FOR_CONSUME if set_all_ready else ProductionStatus.NOT_PRODUCED
 
-    all_fields = [
-        {
-            name: FieldMeta(
-                name=name,
-                dtype=getattr(value, "dtype", None),
-                shape=getattr(value, "shape", None),
+    # unbind nested tensor
+    results: dict = {}
+    for field in tensor_dict.keys():
+        field_data = tensor_dict[field]
+        if batch_size > 1 and isinstance(field_data, Tensor) and field_data.is_nested:
+            results[field] = field_data.unbind()
+        else:
+            results[field] = field_data
+
+    all_fields = []
+    for idx in range(batch_size):
+        dict_of_field_meta = {}
+        for field_name in results.keys():
+            dict_of_field_meta[field_name] = FieldMeta(
+                name=field_name,
+                dtype=getattr(results[field_name][idx], "dtype", None),
+                shape=getattr(results[field_name][idx], "shape", None),
                 production_status=production_status,
             )
-            for name, value in tensor_dict[idx].items()
-        }
-        for idx in range(batch_size)
-    ]
+        all_fields.append(dict_of_field_meta)
 
     return all_fields
 

transfer_queue/storage/managers/base.py

@@ -394,7 +394,14 @@ def _generate_values(data: TensorDict) -> list[Tensor]:
             list[Tensor]: Flattened list of tensors, e.g.,
                           [data[field_a][0], data[field_a][1], data[field_a][2], ..., data[field_b][0], ...]
         """
-        return [row_data for field in sorted(data.keys()) for row_data in data[field]]
+        results: list[Tensor] = []
+        for field in sorted(data.keys()):
+            field_data = data[field]
+            if isinstance(field_data, Tensor) and field_data.is_nested:
+                results.extend(field_data.unbind())
+            else:
+                results.extend(field_data)
+        return results
 
     @staticmethod
     def _shutdown_executor(thread_executor: Optional[ThreadPoolExecutor]) -> None:

transfer_queue/storage/managers/simple_backend_manager.py

@@ -27,6 +27,7 @@
 import zmq
 from omegaconf import DictConfig
 from tensordict import NonTensorStack, TensorDict
+from torch import Tensor
 
 from transfer_queue.metadata import BatchMeta
 from transfer_queue.storage.managers.base import TransferQueueStorageManager
@@ -201,10 +202,21 @@ async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None:
             metadata, self.global_index_storage_unit_mapping, self.global_index_local_index_mapping
         )
 
+        # unbind nested tensor
+        results: dict = {}
+        for field in data.keys():
+            field_data = data[field]
+            if data.batch_size[0] > 1 and isinstance(field_data, Tensor) and field_data.is_nested:
+                results[field] = field_data.unbind()
+            else:
+                results[field] = field_data
+
         # send data to each storage unit
         tasks = [
             self._put_to_single_storage_unit(
-                meta_group.get_local_indexes(), _filter_storage_data(meta_group, data), target_storage_unit=storage_id
+                meta_group.get_local_indexes(),
+                _filter_storage_data(meta_group, results),
+                target_storage_unit=storage_id,
             )
             for storage_id, meta_group in storage_meta_groups.items()
         ]
@@ -221,8 +233,8 @@ async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None:
             per_field_shapes[global_idx] = {}
 
         # For each field, extract dtype and shape for each sample
-        for field in data.keys():
-            for i, data_item in enumerate(data[field]):
+        for field in results.keys():
+            for i, data_item in enumerate(results[field]):
                 global_idx = metadata.global_indexes[i]
                 per_field_dtypes[global_idx][field] = data_item.dtype if hasattr(data_item, "dtype") else None
                 per_field_shapes[global_idx][field] = data_item.shape if hasattr(data_item, "shape") else None
@@ -234,7 +246,7 @@ async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None:
 
         # notify controller that new data is ready
         await self.notify_data_update(
-            partition_id, list(data.keys()), metadata.global_indexes, per_field_dtypes, per_field_shapes
+            partition_id, list(results.keys()), metadata.global_indexes, per_field_dtypes, per_field_shapes
         )
 
     @dynamic_storage_manager_socket(socket_name="put_get_socket")
@@ -432,20 +444,20 @@ def close(self) -> None:
         super().close()
 
 
-def _filter_storage_data(storage_meta_group: StorageMetaGroup, data: TensorDict) -> dict[str, Any]:
-    """Filter batch-aligned data from a TensorDict using batch indexes from a StorageMetaGroup.
+def _filter_storage_data(storage_meta_group: StorageMetaGroup, data: dict) -> dict[str, Any]:
+    """Filter batch-aligned data from a dict using batch indexes from a StorageMetaGroup.
     This helper extracts a subset of items from each field in ``data`` according to the
     batch indexes stored in ``storage_meta_group``. The same indexes are applied to every
-    field in the input ``TensorDict`` so that the returned samples remain aligned across
+    field in the input dict so that the returned samples remain aligned across
     fields.
 
     Args:
         storage_meta_group: A :class:`StorageMetaGroup` instance that provides
             a sequence of batch indexes via :meth:`get_batch_indexes`. Each index
             refers to a position along the batch dimension of the tensors stored
             in ``data``.
-        data: A :class:`tensordict.TensorDict` containing batched data fields. All
-            fields are expected to be indexable by the batch indexes returned by
+        data: A dict containing batched data fields. All fields are expected to
+            be indexable by the batch indexes returned by
             ``storage_meta_group.get_batch_indexes()``.
     Returns:
         dict[str, Any]: A dictionary mapping each field name in ``data`` to a list
@@ -461,7 +473,9 @@ def _filter_storage_data(storage_meta_group: StorageMetaGroup, data: TensorDict)
         return results
 
     for fname in data.keys():
-        result = itemgetter(*batch_indexes)(data[fname])
+        field_data = data[fname]
+        result = itemgetter(*batch_indexes)(field_data)
+
         if not isinstance(result, tuple):
             result = (result,)
         results[fname] = list(result)

tutorial/04_understanding_controller.py

@@ -69,7 +69,7 @@ def demonstrate_partition_isolation():
     train_data = TensorDict(
         {
             "input_ids": torch.tensor([[1, 2, 3], [4, 5, 6]]),
-            "labels": torch.tensor([0, 1]),
+            "labels": torch.tensor([[0], [1]]),
         },
         batch_size=2,
     )
@@ -81,7 +81,7 @@ def demonstrate_partition_isolation():
     val_data = TensorDict(
         {
             "input_ids": torch.tensor([[7, 8, 9], [10, 11, 12]]),
-            "labels": torch.tensor([2, 3]),
+            "labels": torch.tensor([[2], [3]]),
         },
         batch_size=2,
     )