feat(jax): apply training.mixed_precision to the JAX train step
メンテナーはふだん 4 日以内に返信
まだ誰も着手していません。
評価
- 難易度
- 5/5
- 見積もり時間
- 1週間以上
- 初心者へのやさしさ
- 35/100
- issue の種類
- 機能追加
- 明瞭さ
- おおむね明確
- 活発さ
- 静か
- 技術スタック
- python
調査の方向性
まず、JAX trainer における training.mixed_precision の処理と、関連する training および checkpoint のテストを見つけます。正規化された設定がモデル実行、loss、optimizer の更新、reduction、再起動時の動作を通る経路を追跡します。対応モードによって計算 dtype が変更され、未対応の組み合わせが拒否され、数値許容誤差が満たされ、非有限の更新によって checkpoint が破損しないことを確認できれば完了です。
索引モデルが issue の本文から書いたものです。
説明
Summary
The JAX trainer reads training.mixed_precision, but the configuration is not applied to model execution, loss computation, optimizer updates, or checkpoint behavior.
Scope
- Define compute, parameter, output, and reduction dtypes from the normalized mixed-precision configuration.
- Apply the policy consistently to the JAX model and train step.
- Add loss scaling or non-finite update handling when required by the selected low-precision mode.
- Preserve numerically sensitive reductions in an appropriate dtype.
- Serialize enough policy and optimizer state for deterministic restart.
- Reject unsupported dtype/device combinations before training.
Acceptance criteria
- Supported mixed-precision configurations change the actual computation dtype.
- Energy, force, loss, and gradient results remain within documented tolerances of full precision.
- Non-finite gradients do not silently corrupt checkpoints.
- Checkpoint/resume preserves the policy and reproduces the next update.
- Tests cover at least one supported low-precision mode, unsupported configurations, and full-precision regression behavior.
Refs #5756.
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
- 主要言語
- Python
- スター
- 2.1k
- フォーク
- 654
- 平均マージ
- 3日 10時間
- マージ済み PR（30日）
- 13
環境構築
このプロジェクトの開発コンテナを、あなたの GitHub アカウントでブラウザ上に起動します。
- Dockerfile・Docker Compose ファイルなし
- プルリクエストのテンプレートなし
- コントリビューションガイドを読む
はじめの一歩
- issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
- 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
- リポジトリをフォークし、ブランチを切って変更します。
- issue 番号を参照したプルリクエストを送ります。
deepmodeling/deepmd-kit のほかの issue
-
難易度 1/5 1時間未満 初心者へのやさしさ 88/100
deepmodeling/deepmd-kit#6038 · コメント 2 件 ·
メンテナーはふだん 4 日以内に返信
-
bug
難易度 2/5 1〜3時間 初心者へのやさしさ 78/100
deepmodeling/deepmd-kit#5991 ·
メンテナーはふだん 4 日以内に返信
-
Docs enhancement
難易度 2/5 1〜3時間 初心者へのやさしさ 68/100
deepmodeling/deepmd-kit#5766 · コメント 1 件 ·
メンテナーはふだん 4 日以内に返信
-
bug
難易度 2/5 1〜3時間 初心者へのやさしさ 72/100
deepmodeling/deepmd-kit#5689 · コメント 2 件 ·
メンテナーはふだん 4 日以内に返信
-
bug
難易度 2/5 1〜3時間 初心者へのやさしさ 78/100
deepmodeling/deepmd-kit#5686 · コメント 1 件 ·
メンテナーはふだん 4 日以内に返信
deepmodeling/deepmd-kit の issue をすべて見る
似ている issue
-
難易度 1/5 1時間未満 初心者へのやさしさ 72/100
letsencrypt/cp-cps#353 ·
-
難易度 2/5 1〜3時間 初心者へのやさしさ 68/100
-
難易度 2/5 1〜3時間 初心者へのやさしさ 84/100
PedestrianDynamics/pyFDS-Evac#394 ·
メンテナーはふだん 1 日以内に返信
-
難易度 2/5 1〜3時間 初心者へのやさしさ 78/100
DOI-USGS/pywatershed#421 ·
-
難易度 2/5 1〜3時間 初心者へのやさしさ 78/100
python-pillow/Pillow#10087 · コメント 1 件 ·
メンテナーはふだん 1 日以内に返信