Repository: pytorch
Issue: Enhance Memory Timeline Export with Detailed Category-wise Data
Labels: oncall: profiler

### 🚀 The feature, motivation and pitch



Summary:
The `export_memory_timeline` function currently does not provide detailed information in the generated JSON file, which makes it difficult to analyze memory usage in a more granular way. This PR proposes a rewritten version of the function that includes more detailed memory timeline data, grouped by memory usage categories.

Key Changes:

1. Grouped Memory Data by Category:
   - Memory usage data is now grouped by categories, and each category contains a list of timestamps and the corresponding memory usage in GB.
   - The structure of the JSON output is updated to reflect these changes, allowing for better clarity and insight into memory usage per category over time.

2. Category Name Handling:
   - The category names are correctly handled to avoid errors during JSON serialization. Category names are extracted from the `_CATEGORY_TO_COLORS` mapping and are ensured to be represented as strings in the final JSON output.

3. Updated JSON Structure:
   - The exported JSON now contains the following fields:
     - `device`: The name of the device (e.g., `cuda:0`).
     - `max_memory_allocated`: The maximum memory allocated, in GB.
     - `max_memory_reserved`: The maximum memory reserved, in GB.
     - `category_memory_timeline`: A dictionary containing memory usage data for each category, with timestamps and memory usage in GB.

Example of Updated JSON Output:
```json
{
    "device": "cuda:0",
    "max_memory_allocated": 1.23,
    "max_memory_reserved": 1.56,
    "category_memory_timeline": {
        "Category1": [
            {"time": 0.0, "memory_GB": 0.1},
            {"time": 1.0, "memory_GB": 0.2},
            {"time": 2.0, "memory_GB": 0.15}
        ],
        "Category2": [
            {"time": 0.0, "memory_GB": 0.05},
            {"time": 1.0, "memory_GB": 0.1},
            {"time": 2.0, "memory_GB": 0.2}
        ]
    }
}
```

Impact:
- The existing functionality remains unchanged. Users can still use the original memory timeline plotting and exporting features.
- The JSON export functionality now includes more detailed and organized memory timeline data for better analysis.


### Alternatives

    def export_memory_timeline(
        self, path, device, figsize=(20, 12), title=None
    ) -> None:
        import numpy as np
        import json
        mt = self._coalesce_timeline(device)
        times, sizes = np.array(mt[0]), np.array(mt[1])
        if sizes.ndim == 1:
            stacked = np.cumsum(sizes) / 1024**3 
        else:
            stacked = np.cumsum(sizes, axis=1) / 1024**3  
        max_memory_allocated = torch.cuda.max_memory_allocated(device)
        max_memory_reserved = torch.cuda.max_memory_reserved(device)
        category_memory_timeline = {}
        names = ["Unknown" if i is None else i.name for i in _CATEGORY_TO_COLORS]
        num = 0
        for category, color in _CATEGORY_TO_COLORS.items():
            name = names[num]
            num += 1
            i = _CATEGORY_TO_INDEX.get(category)
            
            if i is None:
                continue  
            memory_usage = stacked[:, i] if stacked.ndim > 1 else stacked
            category_memory_timeline[name] = [
                {"time": t, "memory_GB": mem} for t, mem in zip(times / 1e3, memory_usage)  
            ]
    
        timeline_data = {
            "device": device,
            "max_memory_allocated": max_memory_allocated / 1024**3,  # GB
            "max_memory_reserved": max_memory_reserved / 1024**3,  # GB
            "category_memory_timeline": category_memory_timeline
        }

        with open(path, "w") as f:
            json.dump(timeline_data, f, indent=4)

### Additional context

No
```[tasklist]
### Tasks
```


cc @robieta @chaekit @guotuofeng @guyang3532 @dzhulgakov @davidberard98 @briancoutinho @sraikund16 @sanrise