Repository: Lightning-AI/pytorch-lightning
File: src/lightning/pytorch/callbacks/device_stats_monitor.py (Python)

Code before:

        - General Memory Usage:

            - ``allocated_bytes.all.current`` — Current allocated GPU memory
            - ``allocated_bytes.all.peak`` — Peak allocated GPU memory
            - ``reserved_bytes.all.current`` — Current reserved GPU memory (allocated + cached)
            - ``reserved_bytes.all.peak`` — Peak reserved GPU memory
            - ``active_bytes.all.current`` — Current GPU memory in active use
            - ``active_bytes.all.peak`` — Peak GPU memory in active use
            - ``inactive_split_bytes.all.current`` — Memory in inactive, splittable blocks

        - Allocator Pool Statistics* (for ``small_pool`` and ``large_pool``):

            - ``allocated_bytes.{pool_type}.current`` / ``allocated_bytes.{pool_type}.peak``
            - ``reserved_bytes.{pool_type}.current`` / ``reserved_bytes.{pool_type}.peak``
            - ``active_bytes.{pool_type}.current`` / ``active_bytes.{pool_type}.peak``

        - Allocator Events:

            - ``num_ooms`` — Cumulative out-of-memory errors
            - ``num_alloc_retries`` — Number of allocation retries
            - ``num_device_alloc`` — Number of device allocations
            - ``num_device_free`` — Number of device deallocations

        For a full list of CUDA memory stats, see the
        `PyTorch documentation <https://pytorch.org/docs/stable/generated/torch.cuda.memory_stats.html>`_.

    - TPU (via ``torch_xla``)

        - *Memory Metrics* (per device, e.g., ``xla:0``):

            - ``memory.free.xla:0`` — Free HBM memory (MB)
            - ``memory.used.xla:0`` — Used HBM memory (MB)
            - ``memory.percent.xla:0`` — Percentage of HBM memory used (%)

        - *XLA Operation Counters*:

            - ``CachedCompile.xla``
            - ``CreateXlaTensor.xla``
            - ``DeviceDataCacheMiss.xla``
            - ``UncachedCompile.xla``
            - ``xla::add.xla``, ``xla::addmm.xla``, etc.

        These counters can be retrieved using: ``torch_xla.debug.metrics.counter_names()``

    Args:
        cpu_stats: if ``None``, it will log CPU stats only if the accelerator is CPU.
            If ``True``, it will log CPU stats regardless of the accelerator.
            If ``False``, it will not log CPU stats regardless of the accelerator.

    Raises:

Review comment (suggestion):
```suggestion
        `PyTorch documentation <https://docs.pytorch.org/docs/stable//generated/torch.cuda.device_memory_used.html>`_.
```

Code after:

        - General Memory Usage:

            - ``allocated_bytes.all.current`` — Current allocated GPU memory
            - ``allocated_bytes.all.peak`` — Peak allocated GPU memory
            - ``reserved_bytes.all.current`` — Current reserved GPU memory (allocated + cached)
            - ``reserved_bytes.all.peak`` — Peak reserved GPU memory
            - ``active_bytes.all.current`` — Current GPU memory in active use
            - ``active_bytes.all.peak`` — Peak GPU memory in active use
            - ``inactive_split_bytes.all.current`` — Memory in inactive, splittable blocks

        - Allocator Pool Statistics* (for ``small_pool`` and ``large_pool``):

            - ``allocated_bytes.{pool_type}.current`` / ``allocated_bytes.{pool_type}.peak``
            - ``reserved_bytes.{pool_type}.current`` / ``reserved_bytes.{pool_type}.peak``
            - ``active_bytes.{pool_type}.current`` / ``active_bytes.{pool_type}.peak``

        - Allocator Events:

            - ``num_ooms`` — Cumulative out-of-memory errors
            - ``num_alloc_retries`` — Number of allocation retries
            - ``num_device_alloc`` — Number of device allocations
            - ``num_device_free`` — Number of device deallocations

        For a full list of CUDA memory stats, see the
        `PyTorch documentation <https://docs.pytorch.org/docs/stable//generated/torch.cuda.device_memory_used.html>`_.

    - TPU (via ``torch_xla``)

        - *Memory Metrics* (per device, e.g., ``xla:0``):

            - ``memory.free.xla:0`` — Free HBM memory (MB)
            - ``memory.used.xla:0`` — Used HBM memory (MB)
            - ``memory.percent.xla:0`` — Percentage of HBM memory used (%)

        - *XLA Operation Counters*:

            - ``CachedCompile.xla``
            - ``CreateXlaTensor.xla``
            - ``DeviceDataCacheMiss.xla``
            - ``UncachedCompile.xla``
            - ``xla::add.xla``, ``xla::addmm.xla``, etc.

        These counters can be retrieved using: ``torch_xla.debug.metrics.counter_names()``

    Args:
        cpu_stats: if ``None``, it will log CPU stats only if the accelerator is CPU.
            If ``True``, it will log CPU stats regardless of the accelerator.
            If ``False``, it will not log CPU stats regardless of the accelerator.

    Raises: