deepmodeling / deepmodeling/deepmd-kit
feat(pt_expt): add torch profiler lifecycle and trace export
- 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
Assessment
This issue has not been assessed yet.