deepmodeling / deepmodeling/deepmd-kit
feat(jax): implement TensorBoard training output
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Summary
The common schema exposes `training.tensorboard*` options, and the JAX trainer stores related configuration, but training does not currently emit TensorBoard events.
## Scope
- Implement writer creation, metric emission, flushing, and cleanup.
- Honor the enable flag, frequency, and log-directory options independently of console display frequency.
- Emit stable tags for learning rate, training/validation metrics, timing, and task identity.
- Ensure only the designated rank writes a shared event stream.
- Avoid per-step host transfers or imports when no event is due or TensorBoard is disabled.
- Continue global step numbering correctly after restart.
## Acceptance criteria
- No event file is created when TensorBoard is disabled.
- Enabled training emits scalars at the configured frequency and global step.
- Multi-task metrics use unambiguous stable tag names.
- Distributed training has a documented chief-only or rank-separated policy without collisions.
- Writers flush and close on normal and exceptional exits.
- Tests cover enablement, frequency, restart, and cleanup.
Refs #5756.
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Contributor guide
Assessment
This issue has not been assessed yet.