8000 Port to AMD Strix Halo (ROCm), plus portability, safety and tests by gt50 · Pull Request #1 · gt50/mini-AGI · GitHub
[go: up one dir, main page]

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
name: tests

on:
push:
pull_request:

jobs:
test:
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, windows-latest]
python: ["3.11"]
runs-on: ${{ matrix.os }}
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python }}
- name: Install CPU torch
run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
- name: Install the rest
run: python -m pip install -e ".[serve,dev]"
- name: Tests
run: python -m pytest -q
11 changes: 10 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,19 @@ runs/*
__pycache__/
*.pyc

# the training corpus - 1.2 GB of text, regenerated by tools/expand_corpus.py
# the training corpus - 1.2 GB of text, regenerated by corpora/expand.py
data/
runs/corpus_index.json

# secrets
.env
.env.*

# local environments and tool state
.venv*/
venv/
.claude/
.vscode/
.idea/
.pytest_cache/
*.egg-info/
39 changes: 38 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -424,7 +424,7 @@ The right panel shows which subjects are still moving. reasoning, code, stories,

## Running it

1. Make sure you have a CUDA-capable GPU with at least 8 GB of VRAM, and Python 3.10 or newer. The reference machine is an RTX 3070 Laptop GPU with 8 GB.
1. Make sure you have a GPU with at least 8 GB of memory - a CUDA card, or an AMD GPU with a ROCm build of PyTorch (see [AMD and Strix Halo](#amd-and-strix-halo)) - and Python 3.10 or newer. The reference machine is an RTX 3070 Laptop GPU with 8 GB.
2. Clone the repository:
```bash
git clone <repository-url>
Expand All @@ -438,6 +438,8 @@ The right panel shows which subjects are still moving. reasoning, code, stories,
pip install scipy # a few of the analysis tools
```
PyTorch has to match your CUDA version - see [the PyTorch install page](https://pytorch.org/get-started/locally/). The reference environment is torch 2.6.0+cu124 with numpy 1.24.4. Only the first line is needed to train.

Or, with torch already installed for your GPU, `pip install -e .[serve,corpora]` installs everything else from `pyproject.toml`, which leaves torch alone on purpose.
4. Build the corpus. One command downloads the four public datasets and generates the other four lanes:
```bash
python3 -m corpora all # all eight subjects, a few GB
Expand All @@ -457,6 +459,41 @@ The right panel shows which subjects are still moving. reasoning, code, stories,

The run writes a sample log, redraws its graphs as it goes, and checkpoints every few minutes. It is meant to be left alone for days.

### AMD and Strix Halo

ROCm builds of PyTorch answer to the same `torch.cuda` namespace, so nothing in the code is vendor-specific and `--device cuda` is still the flag. What an AMD APU such as the Ryzen AI Max (Strix Halo: Radeon 8060S, gfx1151, up to 128 GB of shared LPDDR5X) needs is a profile, because the hardware differs in ways the code cannot see:

```bash
python train.py --config configs/strix_halo.yaml read --save # continue the released model
python train.py --config configs/strix_halo_large.yaml read --save # a new, larger one
python serve.py --config configs/strix_halo.yaml
python tools/bench_hw.py --config configs/strix_halo_large.yaml # measure before committing to a run
```

`--config` (or `MINAGI_CONFIG`) replaces `config.yaml` for every reader in the process. Both profiles set a `hardware:` block, read by [`minagi/device.py`](minagi/device.py):

| key | what it does |
|---|---|
| `gpu_mem_gb` | the budget growth's ROOM brake measures against. **Set it to what the GPU may really use** - the Variable Graphics Memory in AMD Software on Windows, the GTT limit on Linux. torch reports the firmware carve-out as the card's size, which on an APU says little |
| `ram_tier_on_gpu` | holds the expert pool's RAM tier as device tensors, so paging an expert onto the card is a device-to-device copy rather than a blocking host transfer. System memory *is* the card's memory here |
| `rocm_aotriton` | `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` - the fast SDPA kernels on RDNA 3.5. Without them the cached forward falls back to math attention, which builds the full score matrix at every row |
| `tunableop` | `PYTORCH_TUNABLEOP_ENABLED=1` - tunes the expert dispatch's batched matmuls once and keeps the results in `runs/` |
| `expandable_segments` | the allocator setting train.py has always used; turn it off if your ROCm build warns it is unsupported |

Every run prints a `hardware:` line naming the device, backend, whether memory is unified, the budget and where it came from.

**Install.** Stock PyTorch wheels do not cover gfx1151. AMD publishes builds for it, for Linux and Windows, from its own index:

```bash
pip install --index-url https://rocm.nightlies.amd.com/v2/gfx1151/ --pre torch
```

Then `pip install -e .[serve,corpora]` for the rest. A later `pip install -U` can pull a CUDA torch from PyPI over the ROCm one, so re-run the line above after any upgrade. Check with `python -c "import torch; print(torch.version.hip, torch.cuda.get_device_properties(0))"`, which should name gfx1151.

**Memory.** On Linux, a small BIOS carve-out with a large GTT allowance is the usual setup - raise the TTM limit on the kernel command line (`ttm.pages_limit` counts 4 KiB pages: 27648000 is about 105 GiB) and set `gpu_mem_gb` a little below it. On Windows, set Variable Graphics Memory in AMD Software and set `gpu_mem_gb` to match.

**What to expect.** The 8060S is in the same compute class as the reference RTX 3070 Laptop, with less memory bandwidth, so characters per second will be similar or lower - not higher. What it has is memory, roughly twelve times an 8 GB card's, and the profiles spend it on the axis this model is built around: `configs/strix_halo.yaml` keeps the released model's shape and gives it 64 experts per forward instead of 32, the whole pool in device memory, an 8k reach and a 100 GB disk ceiling; `configs/strix_halo_large.yaml` starts a new model with twice the experts at birth, four times the resident set and a trunk 768 wide. Raising `resident` changes routing, not only memory - a forward is admitted more of what its text asks for - so compare held-out loss before and after rather than assuming it is free. Its numbers are starting points: run `tools/bench_hw.py` on the machine first.

### Everything else

```bash
Expand Down
15 changes: 15 additions & 0 deletions config.yaml
10000
Original file line number Diff line number Diff line change
Expand Up @@ -195,3 +195,18 @@ plots:
enabled: true # redraw the graphs whenever the sample log grows
last: 100 # millions of characters the recent-progress graph
# covers, counted back from wherever reading has got

hardware:
# What the machine is, where the code cannot tell for itself - read by
# minagi/device.py. Every key is optional; left out, a discrete NVIDIA card
# behaves exactly as it always has. See configs/strix_halo.yaml for an APU.
#
# gpu_mem_gb: 8 # the budget growth's ROOM brake measures against.
# # Unset, the card's reported size. Set it on an APU,
# # where the reported size is the firmware carve-out
# unified: false # CPU and GPU share memory. Detected for gfx1150/1/2
# ram_tier_on_gpu: false # hold the expert RAM tier on the device. Defaults
# # to `unified`
# expandable_segments: true # PYTORCH_{CUDA,HIP}_ALLOC_CONF
# rocm_aotriton: false # TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
# tunableop: false # PYTORCH_TUNABLEOP_ENABLED=1, results in runs/
235 changes: 235 additions & 0 deletions configs/strix_halo.yaml
10BC8
Original file line number Diff line number Diff line change
@@ -0,0 +1,235 @@
# mini-AGI on AMD Strix Halo - CONTINUING THE EXISTING MODEL.
#
# python train.py --config configs/strix_halo.yaml read --save
# python serve.py --config configs/strix_halo.yaml
#
# The model block below is the released model's shape, so this profile
# continues weights/ as they are (including the Hugging Face snapshots). Only
# run-time settings change: more experts per forward, the whole pool held in
# device memory, a longer reach, a larger disk ceiling. For a NEW, larger
# model see configs/strix_halo_large.yaml.
#
# Expect characters per second close to an RTX 3070 Laptop's, not above it:
# the 8060S is in the same compute class with less memory bandwidth. What it
# has is memory, and this profile spends it.
#
# mini-AGI — the settings worth changing.
#
# Read when a model is created and every time it reads. Anything not here is
# either derived or was not worth a knob.

model:
# THIS BLOCK BUILDS A MODEL; IT DOES NOT CHANGE ONE. Down as far as the
# depth rule, these are fixed in the manifest when the weights are created.
# Editing them later does nothing - the shapes would no longer match.
d_model: 512 # width of the trunk
n_head: 8 # attention heads
d_ff: 1408 # the trunk's own feed-forward
n_prelude: 2 # dense blocks before the recurrent core; run once
n_recur: 1 # blocks per row
n_coda: 0 # readout blocks after the last row

# ------------------------------------------------------------------ DEPTH
#
# The model runs one row at a time and decides after each whether to stop.
# A row is n_recur + n_coda block applications.
#
# RESCALING RULE: max_steps, halt_prior, bptt_window and train_steps_mean
# are all counted in ROWS, so all four move together if n_recur or n_coda
# changes what a row is. Missing one is silent.

max_steps: 24 # most rows a character may take
min_steps: 1 # rows it must run before halting may stop it

halt_prior: 0.072 # halting is held to a mean of 1/this rows
halt_thresh: 0.9 # at inference, stop at the first row past this
halt_freeze: true # a halted character stops being computed, in
# training and writing alike; later characters read
# its final state at the deeper rows
ponder_beta: 0.01 # weight on the halting KL

bptt_window: 24 # rows gradient reaches back through. Also the
# activation memory, so it sets peak VRAM

train_steps_mean: 12.8 # depth sampled while training: poisson(this) + 1,
# clipped. 0 runs the ceiling every step

# The window attention reaches back over. Reading starts at context_start
# and grows only when the model is still gaining from the far end of what
# it has; context_end is the ceiling and binds at load and at every growth.
context_start: 2048 # where a fresh model begins
context_end: 8192 # rotary, so the existing model can extend into it

context_step: 1 # characters added per growth
context_grow_every_chars: 100_000 # characters between growths
context_gain_min: 0.015 # stop growing below this gain from depth
context_every_chars: 65536 # characters between asking

pool:
# experts, width, depth and top_k are fixed in the manifest, like the model
# block. resident, ram_cache and capacity_factor are run-time and can be
# changed freely.
experts: 64 # how many exist at the start; it grows from here
width: 2048 # hidden units in one expert
depth: 1 # SwiGLU blocks stacked inside one expert
top_k: 8 # how many one character routes through

# HOW EXPERTS ARE CHOSEN - one rule: the text chooses. Every character
# ranks the whole pool with the router and asks for its top_k, and every
# forward's first pass adds its characters' requests to the vote of its
# text - everything read since position 0. A forward - a training window, a
# chunk of held-out file, a prompt, each character of a reply - may use at
# most `resident` different experts: the ones its text has voted for most,
# and a character whose request was not admitted takes its best admitted
# expert instead. A character being written routes among what its prompt
# and the reply so far asked for, as it would inside a training window.
#
# Three tiers: disk holds every expert, RAM caches the recently wanted,
# VRAM holds what the current forward admitted. Only VRAM is scarce.
resident: 64 # experts one forward may use - TWICE the 8 GB card's.
# Changes routing, not only memory: a forward is
# admitted more of what its text asks for
# holding weights, gradient and Adam's moments
ram_cache: 0 # 0 = every expert stays in the tier, which here
# is device memory: disk is read once per expert

capacity_factor: 1.5 # most of a batch one expert may take, as a multiple
# of its fair share; the rest is dropped. 0 removes

# THE BALANCE TERM - how the whole pool gets used. Left alone the router
# keeps choosing the same experts: reading held-out on its own it put 59 of
# 118 on a card across every subject. This term in the loss charges the
# router, in every forward that trains, for the probability it puts on each
# expert over the whole pool, in proportion to that expert's share of recent
# admissions (the last ~1,000 training forwards). Busy experts get dearer and
# idle ones cheaper, so the router's own rows move toward the idle - and
# reading and writing use them too, not only training. A text is still free
# to want few experts; what gets evened out is use across texts.
#
# The weight sets how hard it leans against learning the text on the experts
# in use. An expert no text admits gets no other gradient, so Adam moves its
# row at the usual pace whatever the weight. 0 turns it off. In a 15-minute
# test 0.01 took the experts in use on held-out from 59 to 114, at a cost of
# 0.023 in held-out loss over those minutes.
balance: 0.001

# THE TEMPERATURE OF EXPERT SELECTION - which experts compute a character,
# NOT which character comes out. Characters are always chosen greedily; any
# randomness in what the model writes comes from here.
#
# 0 threshold selection, the rule as it has run: the card holds the
# 32 experts the text's first row asks for most, and every character
# takes its own top 8 of those
# > 0 drawn: at each row, 8 experts are drawn from the text's
# probabilities (the router's full distribution, summed over every
# character in the context window) in proportion to p^(1/T), until
# the card's 32 are full, and every character draws its own 8 from
# its own probabilities over the card, the same way
#
# 1 draws in proportion to the router's probabilities, 0.5 to their
# squares, 0.2 to their fifth powers. Close to 0 the draws approach each
# row's top 8 - which is not quite 0 itself: a drawn card fills 8 a row, so
# the first row has only 8 experts where 0 gives it all 32.
select_temperature: 0

training:
save_every: 5 # minutes between checkpoints

precision: bf16 # what the forward computes in. Weights, gradients
# and the expert files stay fp32

# THE LEARNING RATE IS NOT SCHEDULED. minagi/plasticity.py watches held-out
# and moves the rate in both directions: easing it down without evidence of
# improvement, stepping it back up when held-out jumps and stays there.
chunk: 2048 # characters between steps within the window
lr: 3.0e-4
trunk_lr_mult: 0.1 # the trunk is shared by every domain and is where
# forgetting happens; it learns slower on purpose
weight_decay: 0.1
clip: 1.0

growth:
# Every `every_chars`, add `k` experts if all five brakes are off:
# ROOM there is disk and VRAM for it
# USED the capacity it already has is being asked for
# EARNING the experts added recently are being asked for
# FITS not too many are already inside their trial
# HONEST train and held-out have not separated
every_chars: 2_000_000 # characters between growth decisions
k: 1 # experts added per decision

max_gap: 0.4 # HONEST: train minus held-out. 0 disables

# USED. Staleness is a share of prune.survival_chars: 0.0 when a forward
# admits an expert, 1.0 when it has gone long enough unadmitted to be
# deleted. Past prune.dying_at it counts as dying.
dying_frac_max: 0.35 # stop growing once this share of the pool is dying

# FITS.
max_in_flight: 100 # experts allowed inside their trial at once

# EARNING. Of the experts past their trial and still young, this share must
# have earned a gate.
keep_ratio_min: 0.35
recent_mult: 4.0 # "still young" is age < this x the trial

# ROOM. An expert is 25.2 MB, so 10 GB is about 397 of them.
max_disk_gb: 100.0 # ~4,000 experts at 25.2 MB
mem_frac: 0.85 # of hardware.gpu_mem_gb; leaves the OS its share

birth_gate: 0.001 # gate a new expert starts at - a scale, not a
# verdict. Small so it cannot shout before it earns

prune:
# USE IT OR LOSE IT. An expert is deleted when nothing has admitted it for
# this long, and it is never deleted inside its first survival_chars.
survival_chars: 100_000_000

dying_at: 0.65 # how far through that window an expert must get,
# unaddressed, to count as dying for growth's brake

data:
train: data/train
held_out: data/val
weights: weights # created on the first run if it is not there

shuffle_seed: 0 # which file and which offset each window comes from

# Each sub-folder of `train` is a subject, and reading rotates between them.
# A window never spans two subjects, but a subject may take several in a row.
passage: 32768 # characters read from one file before moving on

decoding:
# Generation is greedy - no random numbers - so two samples taken at the
# same point in training are the same text. What holds off the loops greedy
# decoding usually falls into is a per-character trace that rises when a
# character is used and decays after.
adapt_strength: 2.5 # logits of suppression for a character just used
adapt_decay: 0.88 # half-life of 5.4 characters; 0 disables it

rep_penalty: 1.0 # the weaker alternative, off at 1.0

plots:
enabled: true # redraw the graphs whenever the sample log grows
last: 100 # millions of characters the recent-progress graph
# covers, counted back from wherever reading has got

hardware:
# AMD Ryzen AI Max (Strix Halo): Radeon 8060S, gfx1151, RDNA 3.5, sharing
# up to 128 GB of LPDDR5X with the CPU at ~256 GB/s.
#
# SET THIS to what the GPU may actually use on your machine:
# Windows - the Variable Graphics Memory set in AMD Software (e.g. 96)
# Linux - the GTT limit the kernel allows (ttm.pages_limit), less
# headroom for the OS; or the BIOS UMA size if you set it large
# torch reports the firmware carve-out as the card's size, which on an APU
# says little about what can be allocated, so growth measures against this.
gpu_mem_gb: 96
unified: true
ram_tier_on_gpu: true # the expert RAM tier lives on the device: a
# page-in is a device-to-device copy
expandable_segments: true # turn off if your ROCm build warns it is
# unsupported (some Windows builds)
rocm_aotriton: true # fast SDPA kernels on RDNA 3.5; without them a
# 4k window builds its whole T x T scores per row
tunableop: true # tune each GEMM shape once, results in runs/
Loading
Loading
0