Torch 2.7.1 fixes the MPS SDPA crash that blocks Anima LoRA training on Mac

Every character LoRA for this blog has been trained on RunPod for a while now, because my earlier attempts at training on a Mac ended in failure. Before trying a Muon optimizer implementation, I wanted to know how far AnimaLoraToolkit + Anima-Base LoRA training can actually get on the Mac I have now. Searching turned up almost no information about training Anima models with MPS.

Test environment

Item	Details
Base model	Anima-Base v1.0 (Cosmos-Predict2 2B family)
Training tool	AnimaLoraToolkit (anima_train.py)
Machine	MacBook Pro (M1 Max, 64GB)
Training config	lora_type: "lora" / rank4 / batch1
Data	79 solo images of Kana (pulled from the 3-character combined LoRA set, .txt captions)

This is the Kana I’m training: brown side ponytail, blue scrunchie, and an ahoge are her identifying features.

First stumbling blocks in setup

Running pip install -r requirements.txt as-is fails to install pillow-jxlpy (for jxl image support) in this Mac’s Python 3.13 environment. It’s never imported in the code, and the supported extensions (.jpg/.jpeg/.png/.webp/.bmp) don’t include .jxl, so I dropped it as an unused dependency and installed the rest.

Running the script then stopped at T5Tokenizer requires the SentencePiece library. sentencepiece / tiktoken / protobuf aren’t in requirements.txt but are needed at runtime.

The torch version also had to be sorted out. My past RunPod runs (see the Anima-Base switch article) used torch 2.5.1, but 2.5.1 isn’t available for Python 3.13 on Mac arm64, so I installed the closest version, 2.6.0.

One more thing: the script itself never anticipated MPS. The device selection in anima_train.py is this single line, with no MPS branch.

device = "cuda" if torch.cuda.is_available() else "cpu"

On CPU a 2B model wouldn’t finish testing in any realistic time, so I added an MPS branch.

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

It crashed on torch 2.6.0

With those fixes in place, the run got through the Transformer (685/688 keys matched), the VAE (194/194), the text encoder, LoRA injection into 316 layers, dataset loading, and VAE latent caching. On the first forward pass, the MPS matmul kernel aborted with an LLVM ERROR of incompatible dimensions / invalid shape (SIGABRT, exit 134). The tensor shapes at the crash site were (1,16,64,128) and (1,8,128,64).

It’s a low-level Metal-side crash rather than a Python exception, so try/except couldn’t catch it.

The cause was a PyTorch MPS/SDPA bug

Searching turned up several PyTorch issues reporting the same symptoms in the MPS implementation of scaled_dot_product_attention. #147443 is about the MPS SDPA passing improperly sized tensors to MetalPerformanceShadersGraph and crashing; the issue thread reports that PYTORCH_ENABLE_MPS_FALLBACK=1 doesn’t work around it. #149132 is about MPS SDPA crashing with GQA (different head counts on the query side and the key/value side); the fix was milestoned for 2.7.0. #163597 is a bug newly introduced in 2.8.0, where the fast SDPA kernel breaks on non-contiguous tensors due to wrong stride assumptions. It doesn’t occur on 2.7.1, and the fix is targeted at 2.9.0.

I’ve run into MPS’s weakness around non-contiguous tensors once before, when ComfyUI’s Upscale broke on Mac MPS, and this crash is the same family of problem. The MPS SDPA implementation has a bug around 2.6.0 that crashes on certain tensor shapes, and a different bug that came in with 2.8.0. 2.7.1 sits in between and hits neither (at least this workload ran fine on it).

All 5 steps ran after upgrading to torch 2.7.1

I upgraded torch from 2.6.0 to 2.7.1 (torchvision to 0.22.1) and re-ran with everything else unchanged.

epoch=0 step=1 loss=0.098597 lr=1.00e-04 speed=0.11 it/s
epoch=0 step=2 loss=0.062920 lr=1.00e-04 speed=0.33 it/s
epoch=0 step=3 loss=0.171651 lr=1.00e-04 speed=0.35 it/s
epoch=0 step=4 loss=0.073384 lr=1.00e-04 speed=0.39 it/s
epoch=0 step=5 loss=0.017360 lr=1.00e-04 speed=0.38 it/s

All 5 steps completed without an exception and the LoRA was saved. lora_up.weight should be zero-initialized, but after saving, the sum of absolute values was well above zero, confirming the weights actually changed.

I also generated before/after images with AnimaLoraToolkit’s built-in sampling.

Before training (LoRA weights zero-initialized, effectively plain Anima-Base), a black-haired girl came out with no artifacts. That’s the base model’s plain inference, so of course it works. But in the past SDXL failures, VAE fp16 overflow corrupting samples was one of the suspected causes, and that kind of environment-level breakage would show up in pre-training samples too. At minimum this shows the MPS generation path isn’t broken (the old article only has samples from epoch 1 onward, so there’s no record of whether it was already broken before training). After 5 steps, the output changed to a girl with different hair and animal ears. Kana’s features (brown side ponytail from the training data) aren’t showing yet, but the constant character collapse from last time isn’t happening either, so it looked safe to keep going.

The past SDXL/kohya-ss failures had a clear breakage pattern: mangled hands and literal “ERROR” text from epoch 1 onward, no matter the parameters. Neither the baseline nor the 5-step sample shows that kind of breakage this time.

A 300-step stability test

Five working steps don’t tell you whether a run survives to the end. Among the past SDXL failures, some runs worked in epoch 1 and then died on the way to epoch 2. So I extended the step count to check whether crashes or output collapse appear mid-run.

I initially tried to keep resolution at 512, but the Kana training images are stored at 1024×1024, and past RunPod training (the 3-character combined LoRA and others) ran at 1024 as well. Staying at 512 would mean downscaling the material and would break comparability with the past measurement (2.3 s/step), so I switched to 1024 and re-ran.

Using all 79 Kana images, the run goes 300 steps (79 images at batch1, so roughly 3.8 epochs), generating sample images every 50 steps to check for early signs of collapse.

Quality isn’t much of a concern this round. In my past character LoRA runs, character features start appearing from ep25 or so, and outputs become judgeable somewhere around ep60–100; 300 steps (~3.8 epochs) is a level I normally wouldn’t even look at. The question isn’t “does a proper Kana come out” but whether crashes or breakage (mangled hands, ERROR text) appear as steps accumulate, which is one stage before any quality judgment.

Steps 0–100

Before training (step 0, baseline), a black-haired girl came out with no artifacts.

At step 50 there was no anatomical breakage, but sexual content appeared despite the safe tag in the prompt. The training data (Kana portrait shots, no exposure tags) contains nothing sexual, so this is likely Anima-Base’s own baseline bias (a tendency to exaggerate body shapes with skimpy outfits, already reported in my earlier testing articles) surfacing through the LoRA’s perturbation (image blurred here due to the suggestive content).

At step 100 it flipped: brown side ponytail + purple scrunchie + school uniform (white blouse + red ribbon + navy skirt). Kana’s features from the training data came out clearly. Getting this far at 79 images / ~1.3 epochs is earlier than my rule of thumb (features start appearing from ep25). The step-50 sexual content looks like a temporary wobble. Early samples swing a lot from one image to the next, and it’s too early to read them as a trend.

At no stage did clear breakage like mangled hands or “ERROR” text appear.

At step 150, hair color and the ahoge matched the reference (Kana’s training images) well. The scrunchie was teal, close to the reference blue but not an exact match. Meanwhile twin tails (the reference has a one-sided ponytail) and unrelated elements like wings and ornate outfits still lingered.

The 1e-4 learning rate is 5x my production setting

Features showing this early at 79 images / ~1.9 epochs (150 steps) is much faster than the rule of thumb from my multi-character LoRA runs (features start appearing from ep25). Reviewing the config, I had left learning_rate: 1.0e-4 as-is.

Until now, solo (keichan solo), 2-character (keikana v2), 3-character (trio), and 4-character (4char) runs all used learning_rate: 2.0e-5. The trio article’s config table even annotates it as “Anima’s official, proven-stable value. Higher is unstable”. Yet this run was going at 5x that rate.

AnimaLoraToolkit’s own docs/training-tips.md lists learning-rate guidance by dataset size.

Small dataset (<100 images) 5e-5 to 1e-4 / medium (100-500) 1e-4 to 2e-4 / large (500+) 1e-4 to 3e-4

79 images falls in the “small dataset” bracket with a 5e-5 to 1e-4 range, so the 1e-4 I used sits at the top of that range and isn’t outside the toolkit’s general guidance. But that guidance looks only at dataset size; the 2e-5 that was actually validated on the Anima-Base architecture comes from a different place. The toolkit’s default lora_type being LoKr (more expressive but needs tuning, per the toolkit’s own docs) also suggests its default learning rate may be tuned for LoKr.

A higher LR means bigger per-step updates, which would explain features appearing to emerge early. The speed seen in steps 50–150 may be about the LR rather than about the Mac. After letting the 300 steps finish, I’ll re-run under identical conditions with only the learning rate set back to 2e-5 and compare convergence speed and stability.

At steps 200 and 250, hair color, ahoge, side ponytail, and the blue scrunchie all came out consistently. The twin-tail confusion from step 150 resolved, and unrelated elements like wings and ornate outfits stopped appearing (only the eye color drifts green-to-blue against the brown reference). No breakage appeared at any point.

I couldn’t pin down why convergence was so fast

Features coming together this much at 79 images / 300 steps (~3.8 epochs) is much earlier than the multi-character rule of thumb I started with (features from ep25). Suspecting I was comparing against the wrong baseline, I checked the side experiment inside the same “keichan v2” article that re-baked kanachan herself on Anima-Base. That run used 70 images / repeats4 / lr2e-5 (280 steps per epoch), and Codex’s evaluation was that “T (trigger only) gives brown hair, ahoge, side ponytail, and blue scrunchie at every epoch with a stable character, and the best is ep150.” The earliest epoch in that grid was ep30, so the smallest confirmed-stable point is ep30, i.e. 8,400 steps.

This Mac test ran at lr1e-4 (5x); taking step 200, where the features came together, and discounting naively by 5x still only makes it “equivalent to ~1,000 steps,” leaving a gap of more than 8x against 8,400 steps. I considered other factors, but none were decisive.

Factor considered	Verdict
Training data homogeneity	I hypothesized “mostly close-ups, simple, so it learns fast”, but actually counting the 79 captions gave full body/standing 38, close-up portrait 24, and profiles/back views/action poses 17. The compositions vary a lot. Hypothesis withdrawn
Floating point or implementation differences	mixed_precision: "fp32" (the most precise setting), and the Muon-swap MuonSwapAdamW is an empty subclass of torch.optim.AdamW, so numerically identical. Neither is the cause
MPS vs CUDA computation differences	Even with the same math and the same fp32, exact numerical parity across backends isn’t guaranteed. I didn’t compare how much the per-step update differs between MPS and CUDA, so this can’t be fully ruled out. Separate axis from how fast it runs

The most likely factor is lax judging criteria: past production runs were scored by Codex (an LLM) over a grid of multiple epochs and prompt formats, while here I’m eyeballing one image per checkpoint at a fixed seed42. The same keichan v2 article records the lesson that “a single-seed artifact (broken hands, weird pose) can’t separate overtraining from a generation lottery”, so today’s apparent match could also fall apart if regenerated across multiple seeds.

Since the cause can’t be narrowed to one thing, I decided to run the lr2e-5 comparison and a multi-seed check next.

Per-step time fluctuated between intervals

Working backwards from the sampling timestamps, training speed wasn’t constant.

Interval	Duration	Avg per step
steps 1-50	13m58s	~16.8s
steps 51-100	20m00s	~24s
steps 101-150	16m35s	~19.9s
steps 151-200	16m21s	~19.6s
steps 201-250	20m23s	~24.5s

It doesn’t degrade monotonically: 24s → 19.9s → 19.6s → 24.5s, alternating between slow and fast intervals. Thermal throttling should basically be monotonic degradation or a sustained plateau, so MPS-side cache/memory reorganization around sampling (model.eval() for a 25-step inference, then back to model.train()) feels like the more plausible story, but I can’t say for sure. The measured CPU 70s°C / GPU 80s°C / fans at full didn’t look extreme enough for thermal throttling.

The average works out to about 21 seconds per step. At that rate, the same 22,050 steps as the 3-character combined LoRA would take 22050×21s ≈ 129 hours ≈ ~5.3 days on this Mac. A solo-character production run (the kanachan Anima-Base re-bake above at ep150 = 42,000 steps) works out to 42000×21s ≈ 245 hours ≈ ~10.2 days. RunPod (RTX 5090) finishes a 3-character-scale run in about 14 hours, so the straight comparison puts the gap at about 9x. Note these numbers are specific to this test’s config (resolution:1024, mixed_precision:"fp32", rank4) and could shrink if bf16 works (the past SDXL failures were around bf16/fp16, so I didn’t try it this time).

300 steps at lr=1e-4 ran to the end

The run exited with code 0 and never crashed across 300 steps. Loss over the full 300-step record: min 0.0068, max 0.109, with the first-10-step average at 0.033 and the last-10-step average at 0.039. It kept oscillating without settling, consistent with large per-step updates at a high LR.

The final sample (step 300) was a full-body standing shot with no breakage. Brown hair, ahoge, side ponytail, and blue scrunchie, plus a consistent school uniform (beige cardigan, white blouse, light-blue ribbon, plaid pleated skirt, navy knee socks, loafers, school bag). Only the eye color still leans green against the brown reference. No broken hands or limbs.

Re-running at lr=2e-5

With 300 steps confirmed to run to completion, I set only learning_rate back to the 2.0e-5 used in all my past articles (Anima’s official recommendation is for rank32 setups, so I can’t call it the proper value for rank4, but it lines up for comparison) and restarted under otherwise identical settings (rank4, resolution1024, fp32, the 79 Kana images). The goals: separate whether lr1e-4’s fast convergence is explained by the high LR or mainly by something else (like the lax judging), and see how far Mac training gets at a production-like LR.

Step 0 (baseline) produced numerically identical stats to the lr1e-4 version (mean=0.0126, std=0.6393). That’s expected, since it’s pre-training and LR-independent, but it doubled as a reproducibility check.

At step 50, elements absent from the training data appeared: swimsuit + cat ears + sunglasses. This mirrors the lr1e-4 version’s step 50, which also had sexual content and unrelated elements like wings: lowering the LR does not remove the early-phase wobble. The early wobble is a separate story from the LR; the trigger word just hasn’t bound yet, so it seems to output whatever it likes.

Step 100 was blonde + cat ears + open shirt (black lace underwear) + pink cardigan, with none of Kana’s features (brown hair, side ponytail, blue scrunchie). Compared to the lr1e-4 version’s step 100, which already had the brown side ponytail, purple scrunchie, and school uniform, the difference is stark. At the production LR (2e-5), 100 steps is still nowhere near converged, so the lr1e-4 version’s speed was likely mostly the LR.

At step 150 the hair returned to brown, but with fox ears, a fantasy-warrior outfit, and a spread-legs composition. Kana’s features (side ponytail, blue scrunchie, school uniform) still aren’t showing. The framing (crotch angle) is strongly sexualized, so this one is blurred.

Animal ears keep recurring: step 50 (cat), step 100 (cat), step 150 (fox). The training data contains no animal ears at all, so this looks like Anima-Base’s own habit. Sampling is pinned to sample_seed: 42 for every checkpoint, so this seed may simply favor the base model’s animal-ear direction every time. The LoRA’s signal isn’t strong enough to override it yet. A different seed would likely have produced different outputs.

One more thing: a 50-step sampling interval (about 0.6 epochs at 79 images / repeats1) is a finer granularity than I ever use. Past production checks looked at samples per epoch, where one epoch is 280 steps (70 images × repeats4) or 588 steps (294 images × repeats2), much coarser than this. Checkpoint-to-checkpoint swings where the outfit or character flips completely may be ordinary noise present in any training run, unnoticed simply because I never look at this granularity. Not Mac-specific or LR-specific, just visible now because the observation got finer. That might be the simplest explanation.

At step 200, brown hair + a one-sided ponytail brought it much closer to the reference. But the fox ears persist, the eyes are green (reference is brown), and the outfit is shrine-maiden style.

Despite the safe tag in every prompt, nearly every checkpoint leans suggestive: swimsuit, open shirt, shrine-maiden outfit. Nothing sexual exists in the training data or the prompt, so this is the same Anima-Base baseline bias described at the lr1e-4 version’s step 50, showing up here just like the animal ears.

Step 250 turned explicitly sexual (fox ears still present).

300 steps at lr=2e-5 didn’t capture the character

This run also exited with code 0 and never crashed. Loss: min 0.0074, max 0.224, first-10-step average 0.034 vs last-10-step average 0.036, oscillating without settling, same as the lr1e-4 version.

The final sample (step 300) was a blonde twin-tail girl in an ornate pirate-style outfit (tricorn hat, epauletted jacket, frilled skirt). Not a single one of Kana’s features (brown hair, side ponytail, blue scrunchie, ahoge) is present. Honestly, isn’t this Hololive’s Houshou Marine?

Same 300 steps, same seed, same dataset: the lr1e-4 version reached a clean full-body standing shot, while the lr2e-5 version stayed a character unrelated to Kana. The only difference is learning_rate. The lr1e-4 version’s fast convergence was pretty much down to the high LR. At the production LR (2e-5), 300 steps (~3.8 epochs) didn’t capture the features, consistent with the kanachan Anima-Base re-bake checked earlier (70 images / repeats4 / lr2e-5, smallest confirmed-stable point ep30 = 8,400 steps).

Wall time was 2h15m for run 1 and 1h26m for run 2 (run 2 being faster looks like the effect of some unexplained late-night slowdown/interference, unrelated to LR). Both land around the same 2-hour order. The elapsed time barely differs, yet the single variable learning_rate completely changes the character that comes out.

Verifying run 1’s final LoRA across multiple seeds

Every sample so far was pinned to sample_seed: 42, so the possibility of a single-seed artifact remained to the end. I loaded run 1’s (lr1e-4) final LoRA (muon_swap_stability.safetensors) via resume_lora and generated with 3 different seeds starting from sample_seed: 7 (7/8/9).

All 3 images consistently show brown hair, ahoge, side ponytail, and the blue scrunchie. The step-300 result isn’t single-seed luck: the features genuinely baked in.

However, the body shape, chest size in particular, was exaggerated in all 3. Anima-Base’s baseline bias, visible consistently since the pre-training samples, shows up here too.

The sampling prompt had been the minimal masterpiece, best quality, score_7, safe, 1girl, solo, kanachan, which also bothered me, so I swapped in the established prompt I use for everyday Kana generation (explicit side ponytail, ahoge, ponytail position, outfit specification) and regenerated the same 3 seeds (7/8/9) with the same LoRA.

All 3 lost the body-shape exaggeration and came out with brown hair, ahoge, side ponytail, the correct blue scrunchie, and a red-tie school uniform (3 color variations), though the outfits are actually all over the place. She’s wearing clothes I never specified, so the bake is still weak there. The minimal prompt lets Anima-Base’s habits through, and the more concretely the prompt specifies each element, the more stably the output lands where aimed — confirmed once again.

The 300-step runs both finished without a crash, and the quality-side breakage that plagued the past SDXL/kohya-ss failures (mangled hands and “ERROR” text) never appeared this time.