Repository navigation
Port to AMD Strix Halo (ROCm), plus portability, safety and tests - #1
Merged
Merged
Conversation
Hardware: - minagi/device.py: one place for backend/arch detection, the unified-memory flag, a configurable GPU memory budget, and env vars torch reads at import (allocator, ROCm AOTriton, TunableOp), driven by a `hardware:` config block - growth's ROOM brake measures against that budget instead of the device's reported size, which on an APU is only the firmware carve-out - the expert RAM tier can live on the device (default on unified memory), so paging onto the card is a device-to-device copy; ram_cache 0 = keep all - serve.py runs its forwards under amp (bf16), as training does - hardware banner at startup of train.py read and serve.py - configs/strix_halo.yaml (continue the released model) and configs/strix_halo_large.yaml (a new, larger run); README section - tools/bench_hw.py: SDPA kernels, expert bmm, a full training step and paging latency at a config's shapes Portability and hygiene: - --config / MINAGI_CONFIG select the config file for every reader - encoding="utf-8" on every text open (cp1252 crashed corpus builds) - plot refresh no longer shells out to bash - npz handles closed so os.replace works on Windows - build_paged/_split_trunk_pool moved to minagi/build.py; the library no longer imports train.py - GradSNR restarts when a watched tensor changes size (crashed on the first growth in a fresh run) - pyproject.toml; .gitignore fixes Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- minagi/lock.py: one writer per weights directory (O_EXCL lock file with pid/host/role, stale-lock takeover). train.py read refuses a held directory; serve.py falls back to read-only instead of learning - serve.py: JSON-only /api/chat, Origin and Host checks (cross-site and DNS-rebinding POSTs were trained into the weights), max_new clamped to 2048, learning off by default when bound to a non-loopback host - minagi/build.py build_adamw: one trunk/pool split for train.py read, live.py and the replication probe; the server now trains routers at the pool rate as training does - torch seed mixed with characters read, so resumes stop replaying one random sequence - corpora/fetch.py keeps the HF datasets cache unless --drop-cache, and honours HF_DATASETS_CACHE / HF_HOME - tests/ (config, device, lock, paged round trip, serve refusals) and a CPU CI workflow on Ubuntu and Windows Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2 of 3 tasks
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
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.
Summary
Ports mini-AGI to AMD Strix Halo (Ryzen AI Max+ 395, Radeon 8060S, gfx1151) on ROCm PyTorch for Windows. It also fixes the portability, safety and hygiene issues found while reviewing the project.
ROCm builds use the
torch.cudanamespace, and the code has no custom kernels, Triton ortorch.compile. So the port is mostly about what the hardware is, not about the code.Hardware / Strix Halo
minagi/device.py: backend and arch detection, a unified-memory flag, and a configurable GPU memory budget (hardware.gpu_mem_gb→MINAGI_GPU_MEM_GB→ reported size). It also sets environment variables torch reads at import: allocator, AOTriton, TunableOp. Ahardware:config block drives all of it.ram_cache: 0keeps the whole pool.serve.pyforwards run underamp(bf16), as training does. They ran fp32 before.configs/strix_halo.yamlcontinues the released model: the whole pool in device memory and a 100 GB disk ceiling.residentandcontext_endstay at the released model's 32 and 4096; see Findings on the GPU below.configs/strix_halo_large.yamlstarts a new, larger run: d_model 768, 128 experts at birth, resident 128.tools/bench_hw.pytimes SDPA kernels, expert batched matmuls, a full training step and paging at a config's shapes.Portability
--config FILE/MINAGI_CONFIGselect the config for every reader in the process.encoding="utf-8"on every textopen(). Without it, building the corpus on Windows crashed on cp1252.bash.os.replaceworks on Windows.build_pagedand_split_trunk_poolmoved intominagi/build.py. The library no longer importstrain.py.Safety and correctness
minagi/lock.py: one writer per weights directory.train.py readrefuses a directory that is in use.serve.pyfalls back to read-only.serve.pyaccepts JSON only, checks Origin and Host (blocks cross-site and DNS-rebinding POSTs, which got trained into the weights), clampsmax_newto 2048, and turns learning off by default when bound to a non-loopback host.build_adamwis on 8000 e shared optimizer builder. Before, the server trained routers at the trunk's 0.1× rate.GradSNRrestarts when a watched tensor changes size. This fixes a crash at the first expert growth in a fresh run; it reproduces onmaintoo.corpora/fetch.pykeeps the HF datasets cache unless--drop-cacheis passed.Hygiene
pyproject.toml. torch is deliberately not listed, so a ROCm build is never replaced.tests/: 18 tests covering config, device, lock, a paged create→train→grow→save→reload round trip, and serve refusals..github/workflows/ci.yml: CPU torch on Windows.Scope: tested on Windows only. Linux-specific setup notes were left out on purpose; the code stays portable, as upstream's was.
Findings on the GPU
Measured with the released weights (
Volotat/mini-AGI-undertrained, 886.8M characters read) on a Ryzen AI Max+ 395:residentis not free. At 64, held-out chat went from 0.60 to 0.92 nats and code from 0.71 to 1.19. The model learned to choose within 32 slots, so the continuation profile keeps 32.context_endis not free. Held-out is read at the ceiling. At 8192 the released weights scored 1.66 nats, against 0.62 at 4096, because positions past 4,096 are untrained. Kept at 4096.Test plan
pytestpasses: 18/18 with CPU torch 2.14, and with ROCm torch 2.12.0a0+rocm7.13 on the 8060Strain.py readcreates, grows, saves and resumes a tiny model;serve.pyanswers/api/chattrain.py readrefuses a directory held by a learningserve.py, and the lock is released at exittools/bench_hw.pyon the 8060S (Windows, ROCm 7.13 nightly):strix_halo.yaml(resident 64, 8k context, the earlier version of the profile): 1,000 char/s, 7.8 GB peakstrix_halo_large.yaml: 530 char/s at 8k context, 17.4 GB peak (expert bmm 18 TFLOP/s fwd)--config configs/strix_halo.yaml, 3 local lanes: arithmetic, code, chat):held-out (60 chunks per subject): 0.6218 ± 0.011 before → 0.6055 ± 0.012 after (−0.016)
1.20M characters read; mean training loss 0.585; 1,255 char/s while reading (1.0k char/s over the session, evaluations and samples included)
1.7 experts paged onto the GPU per chunk, 92% from the in-memory cache
growth 127 → 128 experts at a 9.8 GB peak
checkpoints at step 433,537 and 433,936; final save 3.68 GB; lock released at exit
reload check: the saved directory loads (128 experts, step 433,936) and scores like the released weights on the same 4 chunks per subject:
Standard error is about 0.05 at that size, so only chat is clearly better.
🤖 Generated with Claude Code