Использование памяти CUDA

Создание снимка

Обычный способ записи снимка — включить историю памяти, выполнить наблюдаемый код, а затем сохранить файл с сериализованным снимком:

# enable memory history, which will
# add tracebacks and event history to snapshots
torch.cuda.memory._record_memory_history()

run_your_code()
torch.cuda.memory._dump_snapshot("my_snapshot.pickle")

Использование визуализатора

Откройте https://pytorch.org/memory_viz и перетащите файл с сериализованным снимком в визуализатор. Визуализатор — это приложение на JavaScript, которое работает локально на вашем компьютере. Оно не загружает данные снимка.

Включение закрепленной памяти (памяти хоста)

По умолчанию снимки памяти включают только память устройства CUDA. Чтобы также фиксировать выделения закрепленной памяти CPU (например, тензоров, созданных с помощью pin_memory=True), передайте record_pinned_host_memory=True:

torch.cuda.memory._record_memory_history(record_pinned_host_memory=True)

run_your_code()
snapshot = torch.cuda.memory._snapshot()
# Host allocator data is in snapshot["host_segments"] and snapshot["host_traces"]
torch.cuda.memory._dump_snapshot("my_snapshot.pickle")

Чтобы записывать только закрепленную память хоста и полностью пропустить аллокатор устройства CUDA, используйте record_pinned_host_memory=True вместе с record_cuda=False:

torch.cuda.memory._record_memory_history(record_pinned_host_memory=True, record_cuda=False)

Хронология активной памяти

На хронологии активной памяти показаны все действующие тензоры на выбранном GPU за период, охваченный снимком. Перемещайтесь по графику и изменяйте масштаб, чтобы рассмотреть небольшие выделения. Наведите указатель мыши на выделенные блоки, чтобы увидеть трассировку стека, соответствующую моменту выделения блока, и такие сведения, как его адрес. Ползунок детализации можно настроить так, чтобы отображалось меньше выделений и повышалась производительность при большом объеме данных.

История состояний аллокатора

В истории состояний аллокатора отдельные события аллокатора отображаются на временной шкале слева. Выберите событие на временной шкале, чтобы увидеть наглядную сводку состояния аллокатора в этот момент. В сводке отображаются отдельные сегменты, возвращенные cudaMalloc, и то, как они разделены на блоки отдельных выделений или свободного пространства. Наведите указатель мыши на сегменты и блоки, чтобы увидеть трассировку стека, соответствующую моменту выделения памяти. Наведите указатель на события, чтобы увидеть трассировку стека, соответствующую моменту события, например освобождению тензора. Ошибки нехватки памяти отображаются как события OOM. Изучение состояния памяти во время OOM может помочь понять, почему выделение завершилось ошибкой, несмотря на наличие зарезервированной памяти.

В сведениях о трассировке стека также указывается адрес, по которому было выполнено выделение. Адрес b7f064c000000_0 относится к блоку (b) по адресу 7f064c000000, который является «_0-м» случаем выделения по этому адресу. Эту уникальную строку можно найти на хронологии активной памяти и выполнить поиск в истории активных состояний, чтобы изучить состояние памяти в момент выделения или освобождения тензора.

Определение выделений памяти вне PyTorch

Если вы подозреваете, что память CUDA выделяется вне PyTorch, можно собрать исходные сведения о выделениях CUDA с помощью пакета pynvml и сравнить их с выделениями, зарегистрированными PyTorch.

Чтобы собрать исходные данные об использовании памяти вне PyTorch, используйте device_memory_used()

import torch
device_idx = ...
print(torch.cuda.device_memory_used(device_idx))

Справочник по API снимков

torch.cuda.memory._record_memory_history(enabled='all', context='all', stacks='all', max_entries=9223372036854775807, device=None, clear_history=False, compile_context=False, global_record_annotations=False, skip_actions=None, record_pinned_host_memory=False, record_cuda=True) [исходный код]

Включает запись трассировок стека, связанных с выделениями памяти, чтобы можно было определить, что выделило каждый участок памяти в torch.cuda.memory._snapshot().

Помимо сохранения трассировок стека для каждого текущего выделения и освобождения, функция также включает запись истории всех событий выделения и освобождения.

Используйте torch.cuda.memory._snapshot(), чтобы получить эти сведения, а инструменты из _memory_viz.py — чтобы визуализировать снимки.

Поведение буфера

При включенной записи будет сохраняться до max_entries экземпляров TraceEntry. По умолчанию сбор трассировок Python не ограничен (sys.maxsize), поэтому для длительных или бесконечно работающих задач следует задать разумное ограничение, чтобы избежать чрезмерного расхода памяти. Ожидается, что каждая запись занимает несколько КБ.

В длительных рабочих процессах или при меньших значениях max_entries будут храниться только последние накопленные max_entries записей, то есть новые записи будут заменять старые.

Реализация на C++ для справки при реализации кольцевого буфера:

if (record_history) {
  if (alloc_trace->size() < alloc_trace_max_entries_) {
    alloc_trace->emplace_back(te);
  } else {
    (*alloc_trace)[alloc_trace_next++] = te;
    if (alloc_trace_next == alloc_trace_max_entries_) {
      alloc_trace_next = 0;
    }
  }
}

Влияние на задержку

Сбор трассировок Python выполняется быстро (2 мкс на трассировку), поэтому можно рассмотреть возможность включения этой функции в рабочих задачах, если вы предполагаете, что когда-нибудь потребуется отлаживать проблемы с памятью.

Сбор трассировок C++ также выполняется быстро (около 50 нс на кадр); для многих типичных программ это составляет около 2 мкс на трассировку, однако значение может меняться в зависимости от глубины стека.

параметр enabled:

None — отключить запись истории памяти. “state” — сохранять сведения о текущих выделениях памяти. “all” — дополнительно сохранять историю всех вызовов выделения и освобождения. По умолчанию — «all».

тип enabled:

Literal[None, “state”, “all”], необязательно

параметр context:

None — не записывать трассировки. “state” — записывать трассировки для текущих выделений памяти. “alloc” — дополнительно сохранять трассировки для вызовов выделения. “all” — дополнительно сохранять трассировки для вызовов освобождения. По умолчанию — «all».

тип context:

Literal[None, “state”, “alloc”, “all”], необязательно

параметр stacks:

“python” — включать в трассировки кадры Python, TorchScript и inductor; “all” — дополнительно включать кадры C++. По умолчанию — «all».

тип stacks:

Literal[“python”, “all”], необязательно

параметр max_entries:

Максимальное количество событий выделения и освобождения, сохраняемых в записанной истории: max_entries.

тип max_entries:

int, необязательно

параметр clear_history:

Очищать историю при включении; по умолчанию — False.

тип clear_history:

bool, необязательно

параметр skip_actions:

Список типов действий, которые следует пропускать при записи истории памяти. Это позволяет уменьшить расход памяти, исключив из записи некоторые типы событий. Допустимые типы действий:

• “alloc”: события выделения памяти
• “free_requested”: запросы на освобождение (память помечена для освобождения)
• “free_completed”: завершенные операции освобождения (память действительно освобождена)
• “segment_alloc”: выделение сегмента с помощью cudaMalloc
• “segment_free”: возврат сегмента CUDA с помощью cudaFree
• “oom”: исключения нехватки памяти
• “snapshot”: события создания снимков памяти

Например, чтобы пропустить запись событий free_requested: skip_actions=[“free_requested”]

По умолчанию — None (записываются все действия).

тип skip_actions:

list[str], необязательно

параметр record_pinned_host_memory:

Если True, записывать историю памяти также для аллокатора закрепленной памяти CPU (хоста). Трассировки аллокатора хоста будут находиться в ключах host_segments и host_traces снимка, возвращаемого _snapshot(). По умолчанию — False.

тип record_pinned_host_memory:

bool, необязательно

параметр record_cuda:

Если True, записывать историю памяти для аллокатора устройства CUDA. Чтобы записывать только закрепленную память хоста, задайте False (вместе с record_pinned_host_memory=True). По умолчанию — True.

тип record_cuda:

bool, необязательно

torch.cuda.memory._snapshot(device=None, augment_with_fx_traces=False) [исходный код]

Сохраняет снимок состояния памяти CUDA на момент вызова.

Состояние представлено словарем следующей структуры.

class Snapshot(TypedDict):
    segments: List[Segment]
    device_traces: List[List[TraceEntry]]


class Segment(TypedDict):
    # Segments are memory returned from a cudaMalloc call.
    # The size of reserved memory is the sum of all Segments.
    # Segments are cached and reused for future allocations.
    # If the reuse is smaller than the segment, the segment
    # is split into more than one Block.
    # empty_cache() frees Segments that are entirely inactive.
    address: int
    total_size: int  #  cudaMalloc'd size of segment
    stream: int
    segment_type: Literal["small", "large"]  # 'large' (>1MB)
    segment_pool_id: Tuple[
        int, int
    ]  # id of the memory pool owning this segment
    allocated_size: int  # size of memory in use
    active_size: int  # size of memory in use or in active_awaiting_free state
    blocks: List[Block]


class Block(TypedDict):
    # A piece of memory returned from the allocator, or
    # current cached but inactive.
    size: int
    requested_size: int  # size requested during malloc, may be smaller than
    # size due to rounding
    address: int
    state: Literal[
        "active_allocated",  # used by a tensor
        "active_awaiting_free",  # waiting for another stream to finish using
        # this, then it will become free
        "inactive",
    ]  # free for reuse
    frames: List[Frame]  # stack trace from where the allocation occurred


class Frame(TypedDict):
    filename: str
    line: int
    name: str
    # Optional FX debug fields (present when augment_with_fx_traces=True
    # and the frame corresponds to FX-generated code)
    fx_node_op: str  # FX node operation type (e.g., 'call_function', 'output')
    fx_node_name: str  # FX node name (e.g., 'linear', 'relu_1')
    fx_original_trace: str  # Original model source code stack trace


class TraceEntry(TypedDict):
    # When `torch.cuda.memory._record_memory_history()` is enabled,
    # the snapshot will contain TraceEntry objects that record each
    # action the allocator took.
    action: Literal[
        "alloc"  # memory allocated
        "free_requested",  # the allocated received a call to free memory
        "free_completed",  # the memory that was requested to be freed is now
        # able to be used in future allocation calls
        "segment_alloc",  # the caching allocator ask cudaMalloc for more memory
        # and added it as a segment in its cache
        "segment_free",  # the caching allocator called cudaFree to return memory
        # to cuda possibly trying free up memory to
        # allocate more segments or because empty_caches was called
        "oom",  # the allocator threw an OOM exception. 'size' is
        # the requested number of bytes that did not succeed
        "snapshot",  # the allocator generated a memory snapshot
        # useful to correlate a previously taken
        # snapshot with this trace
        "annotate",  # metadata was attached to a live allocation
        # via _annotate_tensor. 'addr' is the allocation's base
        # address and 'user_metadata' holds the annotation
    ]
    addr: int  # not present for OOM
    frames: List[Frame]
    size: int
    stream: int
    device_free: int  # only present for OOM, the amount of
    # memory cuda still reports to be free
    pool_id: Tuple[int, int]  # id of the memory pool for this entry

Параметры:

• device (Device) – Устройство, для которого нужно сделать снимок. Если значение None, снимок создается для текущего устройства.
• augment_with_fx_traces – Если значение True, дополнить кадры трассировки стека отладочной информацией FX, которая сопоставляет сгенерированный код FX с исходным кодом модели. В объекты Frame добавляются поля fx_node_op, fx_node_name, fx_original_trace и fx_node_info. Значение по умолчанию: False.

Возвращает:

Объект словаря Snapshot

torch.cuda.memory._dump_snapshot(filename='dump_snapshot.pickle', augment_with_fx_traces=False) [исходный код]

Сохраняет сериализованную версию словаря torch.memory._snapshot() в файл.

Этот файл можно открыть в интерактивном просмотрщике снимков на сайте pytorch.org/memory_viz

Размер файлов снимков зависит от max_entries и глубины трассировки стека каждой записи и составляет несколько КБ на запись. Для длительных рабочих процессов с большим значением max_entries размер файла может легко достигать нескольких ГБ.

Параметры:

• filename (str, необязательно) – Имя создаваемого файла. По умолчанию — “dump_snapshot.pickle”.
• augment_with_fx_traces (bool, необязательно) – Если значение True, перед сохранением дополнить снимок отладочной информацией FX. Это позволяет сопоставить трассировки стека сгенерированного кода FX с исходным кодом модели. По умолчанию — False.