feat(jax): apply training.mixed_precision to the JAX train step
I maintainer di solito rispondono entro 4 giorni
Nessuno ha ancora preso questa issue.
Valutazione
- Difficoltà
- 5/5
- Tempo stimato
- Più di una settimana
- Idoneità per principianti
- 35/100
- Tipo di issue
- Funzionalità
- Chiarezza
- Abbastanza chiara
- Stato di attività
- Tranquilla
- Stack tecnologico
- python
- Ambito
- machine-learning
Direzione di ricerca
Inizia individuando la gestione di training.mixed_precision nel trainer JAX e i test pertinenti per il training e i checkpoint. Segui la configurazione normalizzata attraverso l’esecuzione del modello, la loss, gli aggiornamenti dell’ottimizzatore, le riduzioni e il comportamento al riavvio; il lavoro è completo quando le modalità supportate modificano il tipo di dati di calcolo, le combinazioni non supportate vengono rifiutate, le tolleranze numeriche sono rispettate e gli aggiornamenti non finiti non possono corrompere i checkpoint.
Scritto dal modello di indicizzazione a partire dal testo della issue.
Descrizione
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
- Lingua principale
- Python
- Stelle
- 2.1k
- Fork
- 659
- Merge medio
- 5g 23h
- PR unite (30g)
- 6
Preparare l'ambiente
Avvia il container di sviluppo del progetto nel browser, con il tuo account GitHub.
- Nessun Dockerfile né file Docker Compose
- Nessun modello di pull request
- Leggi la guida per i contributori
Come iniziare
- Leggi tutta la issue e poi la guida ai contributi del progetto.
- Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
- Fai un fork del repository e lavora su un branch.
- Apri una pull request che faccia riferimento al numero della issue.
Altre issue di deepmodeling/deepmd-kit
-
Difficoltà 1/5 Meno di un'ora Idoneità per principianti 88/100
deepmodeling/deepmd-kit#6038 · 2 commenti ·
I maintainer di solito rispondono entro 4 giorni
-
Docs enhancement
Difficoltà 2/5 1-3 ore Idoneità per principianti 68/100
deepmodeling/deepmd-kit#5766 · 1 commento ·
I maintainer di solito rispondono entro 4 giorni
-
[Code scan] Freeze Paddle DPA3 parameters with stop_gradientForse già presa @njzjz l’ha presa 95 giorni fa. Apertabug
Difficoltà 2/5 1-3 ore Idoneità per principianti 78/100
deepmodeling/deepmd-kit#5686 · 1 commento ·
I maintainer di solito rispondono entro 4 giorni
-
branch: master lang: py todo comment type: todo
Difficoltà 2/5 1-3 ore Idoneità per principianti 72/100
deepmodeling/deepmd-kit#5480 ·
I maintainer di solito rispondono entro 4 giorni
-
branch: master lang: toml todo comment type: todo
Difficoltà 1/5 Meno di un'ora Idoneità per principianti 72/100
deepmodeling/deepmd-kit#5462 ·
I maintainer di solito rispondono entro 4 giorni
Tutte le issue di deepmodeling/deepmd-kit
Issue simili
-
`lvm_compact` has no self-hosted twin, and `docs/158` §5 says every language-face component mustApertabug sev-3 vm
Difficoltà 2/5 1-3 ore Idoneità per principianti 65/100
-
Difficoltà 2/5 1-3 ore Idoneità per principianti 82/100
sierra-research/tau2-bench#635 ·
I maintainer di solito rispondono entro 1 giorno
-
The add-ssm-model skill requires a Cython simulator; the contributing guide allows Python onesForse già presa @fmuia l’ha presa oggi. Apertadocumentation skill
Difficoltà 2/5 1-3 ore Idoneità per principianti 83/100
lnccbrown/ssm-simulators#406 ·
I maintainer di solito rispondono entro 4 giorni
-
Difficoltà 1/5 Meno di un'ora Idoneità per principianti 95/100
btclib-org/btclib-ecc#309 ·
I maintainer di solito rispondono entro 1 giorno
-
Difficoltà 2/5 1-3 ore Idoneità per principianti 74/100
autowarefoundation/autoware_carla_scenario#281 ·
I maintainer di solito rispondono entro 1 giorno