deepmodeling / deepmodeling/deepmd-kit

feat(tf2): implement training profiler support

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

TF2 currently raises `NotImplementedError` when `training.profiling` or `training.enable_profiler` is configured, although both options are exposed by the common schema.

## Scope

- Define and implement the TF2 semantics of both profiler options.
- Own profiler start, step/range control, trace export, and cleanup through a dedicated lifecycle component.
- Honor the configured output file or directory.
- Keep profiler control outside compiled train-step functions.
- Prevent distributed ranks from writing colliding traces.
- Preserve a zero-overhead disabled path.

## Acceptance criteria

- Each profiler option works independently and in combination.
- Valid traces are written to the configured location.
- Resources are finalized on normal completion and exceptions.
- Distributed output follows a documented chief-only or rank-suffixed policy.
- Tests cover lifecycle ordering, trace creation, disabled behavior, and unchanged numerical results.

Refs #5757.

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.