deepmodeling / deepmodeling/deepmd-kit

feat(pt_expt): add torch profiler lifecycle and trace export

Open
#5,820 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

## Summary

Implement `training.enable_profiler`, `training.profiling`, and `training.profiling_file` in pt_expt through a dedicated profiler observer.

## Root cause

Profiler setup, stepping, trace export, and cleanup are lifecycle concerns, but the current pt_expt trainer has no observer layer and silently ignores the profiler configuration.

## Proposed design

Add a `TorchProfilerObserver` using the common training-observer hooks.

It should own the complete profiler lifecycle:

1. create and enter the profiler at training start;
2. call `profiler.step()` after each optimizer step;
3. export the requested Chrome/TensorBoard traces;
4. exit and release the profiler on completion or failure.

If both profiling modes are enabled, one profiler session should serve both outputs rather than instrumenting the training step twice.

For distributed runs, trace filenames must include the rank unless profiling is explicitly chief-only.

## Performance and correctness requirements

- The disabled path must be a true no-op with negligible overhead.
- Profiler control must remain outside the compiled model graph.
- Do not introduce unconditional synchronization into normal training.
- Preserve the established legacy PT meanings of `enable_profiler` and `profiling`.
- Export valid traces after both normal completion and an early controlled stop.
- Avoid file collisions between ranks.

## Acceptance criteria

- Each option works independently and in combination.
- `profiling_file` is honored for single-process training.
- Distributed trace outputs have deterministic, collision-free rank suffixes.
- `profiler.step()` is called exactly once per completed optimizer step.
- Resources are released when training raises an exception.
- Tests verify disabled overhead, lifecycle ordering, and trace creation.

Refs #5755.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.