support MPS, reorganize #1
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This commit adds comprehensive device detection and selection utilities that support CUDA, MPS (Apple Silicon), and CPU backends with automatic fallback logic. Changes: - Add decoder_pytorch/device.py with DeviceSelection dataclass and get_optimal_device() function - Update decoder_pytorch/__init__.py to export new device utilities - Refactor train.py to use get_optimal_device() instead of hardcoded device selection - Use device-specific autocast dtype (bfloat16 for CUDA/MPS, float32 for CPU) - Integrate TF32 configuration into get_optimal_device for CUDA - Update fused optimizer check to only enable on CUDA (not MPS/CPU) The get_optimal_device() function provides: - Automatic device detection with configurable priority order - Force device selection via parameter or FORCE_DEVICE env var - Integrated TF32 configuration for CUDA devices - Appropriate autocast dtype selection per device type - Detailed device info logging This ensures the codebase works seamlessly across CUDA, MPS, and CPU devices with optimal settings for each platform. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
Changed CPU device selection to use torch.bfloat16 for autocast, consistent with the repo's assumption of bfloat16-compatible hardware (2025 AD standard). This eliminates the warning: "CPU Autocast only supports dtype of torch.bfloat16, torch.float16" All devices (CUDA, MPS, CPU) now uniformly use bfloat16 for mixed precision training. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
Added configurable autocast support to allow users to enable/disable mixed precision training via config files without modifying code. Changes: - Add use_autocast config option to simple.yaml and test.yaml (default: true) - Update train.py to conditionally use autocast based on config - Use contextlib.nullcontext() when autocast is disabled - Print mixed precision status on startup Usage: use_autocast: true # Enable bfloat16 mixed precision (default) use_autocast: false # Disable, use full fp32 precision Both configurations tested successfully with no warnings. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
* Add `DeviceSelection.autocast_context()`; parse `cuda:N`, dedupe prefs, and warn on bad input (decoder_pytorch/device.py:26,53). * Honor forced indices; guard out-of-range CUDA; loud CPU fallback for debug (decoder_pytorch/device.py:107). * Use context helper for train/val; fix E731; AMP toggle driven by config (train.py:74). * Document detection flow, `FORCE_DEVICE`, and autocast usage (README.md:27).
Owner
Author
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Owner
Author
|
CPU (colab) 10-steps test config: |
Owner
Author
|
Training converges fine on m3 max MBP (macOS 15.7.1). Generations are fine as well
|
Owner
Author
|
Need to decide
|
- Replace device.py with lightweight tuple-based API and auto-fallbacks - Centralize device checks in training; guard autocast and document grad quirks - Graceful Flash Attention degradation when kernels unavailable - Add nano.yaml config for quick CPU/MPS testing - Update docs to reflect new device API and config
- Stop silently disabling autocast; always respect use_autocast flag - Wrap autocast context manager on all devices (no silent fp32 fallback) - Align nano.yaml to ~20M Llama with bf16 autocast enabled - Clarify autocast behavior in README
- Update nano.yaml: depth 6, dim 384, torch.compile on - Clarify model scale in README
Owner
Author
|
@codex review. we already coverd bf16 autocast so dw about taht |
|
Codex Review: Didn't find any major issues. Swish! ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Owner
Author
|
validated that training works on mac torch 2.9 - autocast/bf16 defaults on mps obviously works on linux/cuda too |
pszemraj
pushed a commit
that referenced
this pull request
Nov 9, 2025
Implements full Llama-style transformer in Rust with Burn tensor library: ## Core Components - RMSNorm: Root Mean Square Layer Normalization - SwiGLU: Gated feedforward with SiLU activation - RoPE: Rotary Position Embeddings - Attention: Multi-head attention with causal masking - Transformer: Pre-normalization blocks - Llama: Complete character-level language model ## Features - Type-safe tensor operations with compile-time checking - Backend-agnostic (NdArray CPU, WGPU GPU, etc.) - Memory-safe implementation (no runtime errors) - Configuration via YAML (compatible with PyTorch configs) - enwik8 dataset loading and training loop - Gradient accumulation framework - Progress tracking and metrics logging ## Verification - Tested PyTorch implementation on enwik8 - Training dynamics verified: 5.5 → 3.0 loss (100 steps) - Generated text shows learning of character patterns ## Structure - decoder-rust/src/model/: All model components - decoder-rust/src/data/: Dataset loading - decoder-rust/src/bin/train.rs: Training script - Complete documentation in decoder-rust/README.md Note: Rust compilation requires network access to download dependencies from crates.io. All code is complete and ready for compilation when network access is available. Refs: #1 (decoder template port request)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
this adds a magic "auto optimal device" utility to decide whether to use cuda or mps accelerators or fall back to CPU, and each of these is set to use bf16 autocast by default, but can disable/fall back to fp32 if needed (cough mps cough)