setup_torch_profiler¶

torchtune.training.setup_torch_profiler(enabled: bool = False, cpu: bool = True, cuda: bool = True, profile_memory: bool = False, with_stack: bool = False, record_shapes: bool = True, with_flops: bool = False, wait_steps: Optional[int] = None, warmup_steps: Optional[int] = None, active_steps: Optional[int] = None, num_cycles: Optional[int] = None, output_dir: Optional[str] = None) → Tuple[profile, DictConfig][source]¶

设置profile并返回带有设置后更新的分析器配置。

分析器配置可以在profiler键下的配置中提供，其布局如下：

profiler:
  _component_: torchtune.training.setup_torch_profiler
  enabled: bool
  # Output directory of trace artifacts
  output_dir: str

  # torch.profiler.ProfilerActivity types to trace
  cpu: bool
  cuda: bool

  # Trace options
  profile_memory: bool
  with_stack: bool
  record_shapes: bool
  with_flops: bool

  # torch.profiler.schedule args
  wait_steps: int
  warmup_steps: int
  active_steps: int
  num_cycles: int

分析器的更新计划与优化器步骤相关（例如，如果 gradient_accumulation = 2，那么分析器将每2个批次更新一次）。

如果配置缺少选项，将选择合理的默认值：

如果没有指定活动，分析器将默认使用CPU + CUDA
如果未指定计划，分析器将默认使用 DEFAULT_SCHEDULE
某些选项将被覆盖（with_stack 和 record_shapes），具体取决于其他选项的要求（例如，profile_memory 需要 with_stack 和 record_shapes）。

注意

启用分析器将导致训练速度降低。
设置 profile_memory: True 将会生成大型的跟踪文件。
分析器的调度是上下文相关的。在每个批次迭代中调用profiler.step()但在梯度累积范围之外，将会在每个前向/后向步骤中step分析器。在每个批次迭代中调用profiler.step()但在梯度累积范围之内，将会在每个优化器更新步骤中step分析器，使得每个step包含多个前向/后向传递。

Parameters:

enabled (bool) – 启用pytorch分析器。默认为False。
cpu (bool) – 启用CPU性能分析。默认值为True。
cuda (bool) – 启用cuda性能分析。默认值为True。
profile_memory (bool) – 分析内存使用情况。默认为 False。
with_stack (bool) – 分析堆栈。默认值为 False。
record_shapes (bool) – 记录形状。默认值为 True。
with_flops (bool) – 分析浮点运算次数。默认值为 False。
wait_steps (可选[int]) – 等待的步数。映射到 wait 的 torch.profiler.schedule 参数。
warmup_steps (可选[int]) – 预热步数。映射到 warmup 的 torch.profiler.schedule 参数。
active_steps (可选[int]) – 活动时间的步数。映射到 active 的 torch.profiler.schedule 参数。
num_cycles (可选[int]) – 分析循环的次数。映射到 repeat 的 torch.profiler.schedule 参数。
output_dir (可选[str]) – 跟踪文件输出路径。

Returns:

元组[torch.profiler.profile, DictConfig]