deepmodeling / deepmodeling/deepmd-kit

feat(pt_expt): add rank-aware TensorBoard observation

Open
#5,821 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.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

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.