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]

设置 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 都包含多个前向/后向传递。

参数:

• 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) – 分析 flops。默认为 False。

• wait_steps (Optional[int]) – 等待时间（步数）。映射到 torch.profiler.schedule 的 wait kwarg。

• warmup_steps (Optional[int]) – 预热时间（步数）。映射到 torch.profiler.schedule 的 warmup kwarg。

• active_steps (Optional[int]) – 活动时间（步数）。映射到 torch.profiler.schedule 的 active kwarg。

• num_cycles (Optional[int]) – 分析周期数。映射到 torch.profiler.schedule 的 repeat kwarg。

• output_dir (Optional[str]) – 跟踪文件输出路径。

返回:

Tuple[torch.profiler.profile, DictConfig]