Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

feat(jax): consume the common optimizer configuration

Đang mở
  • Có gì không ổn?
#5,867 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 4 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
48/100
Loại issue
Tính năng
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Ít trao đổi
Công nghệ
python
Lĩnh vực
machine-learning

Hướng nghiên cứu

Bắt đầu bằng cách xác định training path của JAX, nơi xây dựng chuỗi Optax Adam cố định và xử lý optimizer block cấp cao nhất đã được chuẩn hóa. Truy vết các test checkpoint và restart hiện có, sau đó bổ sung coverage cho các lựa chọn được hỗ trợ, các cập nhật một bước mang tính tất định, các biến thể không được hỗ trợ và optimizer state được khôi phục. Công việc được xem là hoàn tất khi transformation và các parameter đã chọn được sử dụng mà không ghép việc lập lịch learning rate với việc lựa chọn optimizer.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

enhancement

Summary

JAX training currently constructs a fixed Optax Adam chain and does not consume the normalized top-level optimizer block. Optimizer choices and parameters accepted by the common schema can therefore be ignored.

HybridMuon for DPA4 is tracked separately in #5771; this issue covers the general optimizer configuration and dispatch contract.

Scope

  • Read and validate the normalized optimizer configuration.
  • Add a JAX optimizer factory for the optimizer types that JAX can support.
  • Wire optimizer-specific parameters into the selected Optax transformation.
  • Fail explicitly for unsupported optimizer types instead of silently using Adam.
  • Preserve optimizer state through checkpoints and restarts.
  • Keep learning-rate scheduling separate from optimizer selection.

Acceptance criteria

  • Selecting a supported optimizer changes the transformation actually used by training.
  • Supported optimizer parameters affect deterministic one-step update tests.
  • Unsupported schema variants produce an actionable backend-specific error.
  • Adam's current behavior remains covered by regression tests.
  • Checkpoint/resume restores all optimizer state and reproduces the next update.
  • The factory provides the integration point required by #5771 without duplicating configuration parsing.

Refs #5756.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

Ngôn ngữ chính
Python
Star
2.1k
Fork
661
Merge trung bình
5 ngày 13 giờ
Pull request đã merge (30 ngày)
11

Chuẩn bị môi trường

Mở trong Codespaces

Khởi chạy dev container của dự án ngay trên trình duyệt, bằng tài khoản GitHub của bạn.

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của deepmodeling/deepmd-kit

Tất cả issue của deepmodeling/deepmd-kit

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.