deepmodeling / deepmodeling/deepmd-kit
feat(pt_expt): add rank-aware TensorBoard observation
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Summary
Implement `training.tensorboard`, `training.tensorboard_freq`, and `training.tensorboard_log_dir` in pt_expt through a rank-aware training observer.
## Root cause
The common training loop has no observer lifecycle, and the pt_expt trainer does not consume the schema-exposed TensorBoard options.
## Proposed design
Introduce a lightweight `TrainingObserver` interface with hooks such as:
- `on_train_begin`;
- `on_step_end`;
- `on_display`;
- `on_checkpoint`;
- `on_train_end`.
Implement `TensorBoardObserver` under `deepmd/pt_expt/train` and allow legacy PT to reuse it.
The observer should receive `RankContext`, step metrics, learning rate, timing, and optional gradient diagnostics. Only the chief rank creates and owns a `SummaryWriter`; distributed metrics are reduced only on logging steps.
## Performance and correctness requirements
- No TensorBoard import, writer, filesystem access, or metric synchronization when disabled.
- Honor `tensorboard_freq` independently of `disp_freq`.
- Avoid per-step device-to-host transfers when no TensorBoard event is due.
- Keep all observer operations outside compiled model execution.
- Flush and close the writer on normal completion and exceptional exit.
- Prevent multiple ranks from writing to the same event stream.
## Acceptance criteria
- No event file is created when `tensorboard=false`.
- Exactly the chief rank writes when distributed training is active.
- Scalars are emitted at the configured frequency with the correct global step.
- Learning rate, training metrics, validation metrics, and timing have stable tag names.
- Restarting training continues with restored step numbers without overwriting earlier events.
- Single-process and multi-process tests cover writer creation, frequency, and cleanup.
Refs #5755.
Contributor guide
Assessment
This issue has not been assessed yet.