Apple Silicon support: arm64 build fixes + MPS device for torch backend#609
Apple Silicon support: arm64 build fixes + MPS device for torch backend#609oceanapplications wants to merge 23 commits into
Conversation
- build.sh: gate -mavx2/-mfma on x86_64 (arm64 clang rejects them) - build.sh: macOS OpenMP via -Xpreprocessor -fopenmp with Homebrew's omp.h, linked against torch's bundled libomp.dylib. Linking a second OpenMP runtime (e.g. Homebrew's) into a process that imports torch aborts at startup or segfaults in the vecenv's parallel regions. - torch_pufferl: select mps when available and _C has no CUDA; move actions to host memory before cpu_step; round-trip advantage computation through CPU for non-CUDA accelerators - pufferl: fall back to the torch backend automatically when _C was built with --cpu, so 'puffer train env' works without --slowly Verified on M-series (breakout): 605K SPS on MPS vs 190K on CPU. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
…ckend Shipped configs contain sweep-produced floats (e.g. num_layers = 2.11327 in cartpole.ini). The native backend truncates them on assignment to C ints; the torch backend passed them straight to nn.Linear and crashed with 'float object cannot be interpreted as an integer'. Fixes the torch (--slowly) backend for 17 of the 38 currently-buildable ocean envs, on all platforms. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
ocean env matrix on Apple Silicon (arm64, --cpu build, MPS training)37/38 buildable envs train one full epoch on MPS. 20 build failures: 17 are broken on all platforms on current master (missing ocean/env_binding.h or OBS_TENSOR_T defines), 1 is x86-only SIMD (craftax_classic), and matsci/nethack/impulse_wars need external deps or unvendored headers.
|
The MPS multinomial kernel can intermittently return indices outside [0, num_categories) (pytorch#136623, still unfixed upstream; the fix PR pytorch#170195 was closed unmerged). In long MPS training runs this surfaced as an intermittent 'AcceleratorError: index N is out of bounds' raised at the next sync point — the .cpu() transfer in compute_puff_advantage — with N ~ 2x total_agents, because the bad index from the prioritized-replay multinomial feeds lazily-queued gathers (obs[idx]) and scatters (ratio[idx], val[idx]) that only validate at materialization. Clamp multinomial output on MPS at both call sites: minibatch segment sampling and action sampling (where an out-of-range action would be memcpy'd into the C envs and corrupt memory silently instead of raising). Cost is one elementwise op; out-of-range draws are ~1e-5 rare. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Include paths (pybind11, numpy, sysconfig) and the torch libomp lookup used bare 'python', which fails with ModuleNotFoundError when the target venv is not the active interpreter. Resolve once at the top: PYTHON=$path ./build.sh <env> --cpu now works from any shell. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
f506345 to
62c1220
Compare
|
Used to train this 6 link inverted pendulum on a MacBook. So real world tested over hundreds of billions of steps. demo_250th.trimmed.mp4 |
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0 vecenv.h compile-time binding (OBS_SIZE 104, one discrete head of 4). Env struct gains num_agents/rng and switches actions/terminals to the float pointers vecenv assigns; c_reset now writes initial observations. frameskip default (4, from the old python wrapper) added to the config. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0
vecenv.h compile-time binding: OBS_SIZE 1 (ByteTensor), one discrete
head of 2 (ACT_SIZES {2}). Env struct gains num_agents/rng and switches
actions/terminals to the float pointers vecenv assigns. No config keys
needed beyond the existing size.
Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0
vecenv.h compile-time binding: OBS_SIZE 256 (4 features x NUM_BOIDS=64,
guarded against config mismatch at init), ACT_SIZES {5, 5}. Env struct
gains num_agents/rng and switches terminals to the float pointer vecenv
assigns. The (action - 2)/4 scaling the deleted python wrapper applied
moves into c_step so raw discrete actions keep the same velocity deltas.
report_interval added to the config with the old python default (1).
Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0
vecenv.h compile-time binding: OBS_SIZE 100 (3*NUM_ARMIES=2 + 4*16 + 22
+ 8, guarded against config mismatch at init), ACT_SIZES {1, 1, 1}
(continuous 3-vector, was Box(-1, 1, (3,))). num_agents on the struct
now means learner buffer rows as vecenv requires; a new num_entities
field (2x, matching the num_agents*2 the 3.x wrapper passed to C)
covers the scripted opponent half, so the old num_agents/num_agents/2
loops become num_entities/num_agents. terminals switches to the float
pointer vecenv assigns, rng added, the stale compute_observations
stride assert corrected, and attack_aa gains its missing return false
(now an error under -Werror=return-type). width/height added to the
config with the old python defaults (1920/1080).
Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0 vecenv.h compile-time binding: OBS_SIZE 64 (8x8 board, uint8/ByteTensor), one discrete head of 512 (size*size*8 move types). Env struct gains num_agents/rng and switches actions/terminals to the float pointers vecenv assigns; standalone demo checkers.c updated to match. No config keys added (size=8 already present). Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0
vecenv.h compile-time binding: OBS_SIZE 1 (float32/FloatTensor), one
discrete head of 2. Env struct gains num_agents/rng and switches
actions/terminals to the float pointers vecenv assigns; standalone demo
memory.c updated to match. Config: add length=4 (old python default) to
config/memory.ini, and drop the stale "memory" entry from
config/sanity.ini's env_name list (a 3.x leftover whose envs no longer
exist) — it glob-matched before memory.ini in load_config, handing the
env an empty [env] section and crashing my_init's dict_get("length").
Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0 vecenv.h compile-time binding: OBS_SIZE 1 (uint8/ByteTensor, constant zero obs), one discrete head of 2 (left/right). Env struct gains num_agents/rng and switches actions/terminals to the float pointers vecenv assigns; allocate_World updated to match. No config keys added (mean_left/mean_right/var_right already present). Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0
vecenv.h compile-time binding (OBS_SIZE 21, MultiDiscrete heads {9,5}).
Env struct gains rng and switches actions/terminals to the float
pointers vecenv assigns. num_agents already present in config.
Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0 vecenv.h compile-time binding (OBS_SIZE 49 = (2*vision+1)^2 with vision 3, ByteTensor obs, one discrete head of 5). Env struct gains rng and switches actions/terminals to the float pointers vecenv assigns; terminal memsets updated for float width. No cross-env shared state — each env owns its grid, so the default my_vec_init path suffices. Config: num_agents changed from the old python-list [8] to scalar 8 (C dict values are doubles); width/height added with the old python defaults (32). Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0 vecenv.h compile-time binding (OBS_SIZE 10, ByteTensor obs, one discrete head of 4, matching the old python wrapper spec). Env struct gains rng and switches actions/terminals to the float pointers vecenv assigns. Adds the previously missing c_close (frees the render client if open); no config keys needed — my_init reads no kwargs, mirroring the old no-arg my_init. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Rewrite binding.c from the 3.x env_binding.h CPython glue to the 4.0 vecenv.h compile-time binding (OBS_SIZE 4 uint8 obs as ByteTensor, one discrete head of 3). Env struct gains num_agents/rng and switches actions/terminals to the float pointers vecenv assigns; the standalone allocator follows suit. c_reset already writes initial observations. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Replace the stale OBS_TYPE/ACT_TYPE macros with the OBS_TENSOR_T ByteTensor define vecenv.h now expects, switch the actions pointer from double* to the float* vecenv assigns, and add the rng field vecenv seeds per env. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Replace the stale OBS_TYPE/ACT_TYPE macros with the OBS_TENSOR_T FloatTensor define vecenv.h now expects, switch the actions pointer (struct and standalone allocator) from double* to the float* vecenv assigns, and add the rng field vecenv seeds per env. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Replace stale OBS_TYPE/ACT_TYPE defines with OBS_TENSOR_T ByteTensor (env writes unsigned char observations). Add unsigned int rng to CSnake (vecenv seeds envs[i].rng) and switch actions to float* / observations to unsigned char* to match the vecenv buffer types under -Werror=incompatible-pointer-types. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Replace stale OBS_TYPE/ACT_TYPE defines with OBS_TENSOR_T FloatTensor and switch actions to float* to match vecenv buffer types under -Werror=incompatible-pointer-types. Add the [env] keys my_init reads that were missing from config/whisker_racer.ini (llw_ang, flw_ang, frw_ang, rrw_ang, render, i), using the old Python-side defaults; dict_get asserts on missing keys. Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Add unsigned int rng (vecenv seeds envs[i].rng) and the num_goals field binding.c sets to the Scape struct, and switch actions to float* to match the vecenv action buffer under -Werror=incompatible-pointer-types. Add config/scape.ini: [vec]/[train] modeled on pong/snake with a modest total_agents, [env] providing the width/height keys my_init reads (1080x720, matching the old scape.c demo defaults). Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
craftax_classic.h unconditionally included <immintrin.h> and used ~44
AVX-512 intrinsics in the Perlin worldgen block of generate_world(),
which broke compilation on arm64 (Apple Silicon).
Guard structure:
- The <immintrin.h> include and a new CRAFTAX_AVX512 macro are gated on
__AVX512F__ or (x86/x86_64 + GCC/Clang), since build.sh compiles with
only -mavx2 and the intrinsics are enabled per-function via
__attribute__((target("avx512f,..."))), which is now gated the same way.
- Inside generate_world(), a single #if CRAFTAX_AVX512 / #else / #endif
wraps the noise-generation block: the original AVX-512 code on x86, a
lane-for-lane scalar equivalent elsewhere. Each vector lane maps to one
column (c = c_base + lane); _mm512_cvttps_epi32 becomes int truncation,
the fmsub/fmadd polynomial is perlin_interp(fy), and
_mm512_permutexvar_ps(yN, load(&tab[k][rowM])) becomes tab[k][rowM+yN].
Function signature and call sites are unchanged.
Verified: ./build.sh craftax_classic --cpu builds on arm64 and the MPS
smoke test prints OK steps=524288.
Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
nproc does not exist on macOS; the nethack NLE build step failed before even reaching the compiler. (The NLE build itself still fails on macOS inside the vendored NetHack fork's C++ visibility attributes — upstream NLE has no macOS support — but the build script bug is real regardless.) Co-Authored-By: Claude Fable 5 <[email protected]> Claude-Session: https://claude.ai/code/session_019AsyRcLQqeJondzSTM6xsn
Update: 17 previously-broken ocean envs fixed (now 54/58 runnable envs train on MPS)Follow-up to the matrix above. The 20 build failures broke down into three groups; this branch now fixes 17 of them, each verified with an arm64 3.x -> 4.0 binding migrations (12) - these envs still carried vecenv.h drift fixes (5) - benchmark, blastar, snake, whisker_racer, scape had partially-migrated bindings (stale x86-only SIMD (1) - craftax_classic's AVX-512 Perlin worldgen now has a scalar fallback; the vector path is preserved for x86 (it was enabled via per-function Bonus finds while migrating: a stale Still broken (documented, not fixed): impulse_wars (this repo is missing the upstream env's vendored Caveat: Full re-verification on M4 Max, macOS 26.5.1 (all 17 fixed envs + pong/breakout regression checks):
|
Apple Silicon support: arm64 build fixes + MPS device for the torch backend
Makes
./build.sh <env> --cpubuild on Apple Silicon and the PyTorch (--slowly) backend train on the M-series GPU via MPS. Related: #590, #507, #422, #532.Changes
build.sh
-mavx2 -mfmaonx86_64(arm64 clang rejects them; same idea as Conditionally apply AVX2/FMA flags by architecture (fix Apple Silicon build) #590).-Xpreprocessor -fopenmp+ Homebrew'somp.h, but link against torch's bundledlibomp.dylib. This is the subtle one: linking Homebrew's libomp builds fine, but any process that also imports torch then has two OpenMP runtimes loaded — it either aborts at startup (OMP: Error #15) or segfaults insidecpu_vec_step's parallel regions. Falls back to Homebrew's libomp when torch isn't importable at build time.brew install libompwhenomp.his missing.pufferlib/torch_pufferl.py
cuda→mps(when available) →cpu.cpu_step(the vecenv memcpys from the raw pointer; an MPSdata_ptr()segfaults).compute_puff_advantage: round-trip through CPU for non-CUDA accelerators, since_C.puff_advantage_cpureads raw host pointers. (A native Metal kernel like feature: add mps kernel forcompute_puff_advantage#422 would avoid the copy; this keeps the diff minimal.)pufferlib/pufferl.py
_Cwas built with--cpu(nocreate_pufferl), sopuffer train <env>works on macOS without knowing about--slowly.Results (M-series Mac, macOS 25.5 / Apple clang 21, torch 2.12.1)
Benchmark on
breakout(default config, 4096 agents, 32.5K-param policy), steady-state over 3 epochs:Losses match CPU training qualitatively; checkpoints save/load fine.
Also ran a full build + one-epoch MPS training smoke test of every env in
ocean/on arm64: 37 of the 38 currently-buildable envs train on MPS (the 38th,squared_continuous, has no config file). Includes chess, craftax, drive, nmmo3, moba, and terraform. The 20 envs that don't build fail for reasons unrelated to this PR: 17 are broken on all platforms on current master (references to a missingocean/env_binding.h, orbinding.cmissing itsOBS_TENSOR_Tdefine), pluscraftax_classic(x86-only AVX-512),matsci(needs LAMMPS), andnethack/impulse_wars(external/unvendored deps).The second commit fixes a torch-backend crash that predates this PR and affects all platforms: shipped configs contain sweep-produced float values (
num_layers = 2.11327incartpole.ini), which the native backend truncates to C ints but the torch backend passed straight tonn.Linear. 17 of the 38 buildable envs crashed on this before the fix.Not addressed
craftax_classic: hand-written AVX-512 obs path (44_mm512_*intrinsics, no scalar fallback) — x86-only regardless of this PR.ocean/env_binding.h, missingOBS_TENSOR_Tdefines).