deepmodeling / deepmodeling/deepmd-kit

feat(jax): implement TensorBoard training output

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

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

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.