The bug

UnslothTrainer._configure_online_tokenization resolves how many passes a max_steps run makes over the split, and hands that number to the online-tokenization gate, whose only use of it is the epochs > 1.0 veto. It read the process count like this:

world_size = max(1, int(os.environ.get("WORLD_SIZE", "1")))

dataset_bounds.world_size_from_env() already reads eight launcher variables plus the MLX rank files, and both loaders use it for the row bound. Reading one variable here means these launches resolve as single-process and get the lazy re-tokenizing view on a run that makes more than one pass, which is the case the veto exists to keep on the Arrow cache:

• mpirun, which sets OMPI_COMM_WORLD_SIZE and never WORLD_SIZE (Open MPI env vars)
• MPICH and Intel MPI under Hydra, and Slurm's srun with the pmi2 plugin, which set PMI_SIZE (mpich hydra, slurm pmi2)
• an environment that kept torchrun's LOCAL_WORLD_SIZE and lost the global WORLD_SIZE
• mlx.launch's NCCL backend, which is CUDA-only and so does reach this path

This is not an Apple silicon fix and does not claim to be one. worker.py routes Apple silicon to run_mlx_training_process and returns before UnslothTrainer() is ever constructed, so this method is unreachable there. The rank-file branch is here for consistency with worker.py and unsloth_cli/_inference.py, not because that path is reachable today.

A second bug in the same expression

int("auto") and int("") raise. The surrounding try swallows it, the pass count is left unresolved, and an unresolved step-capped run reads as infinite passes, so the gate vetoes. A junk WORLD_SIZE therefore silently disabled online tokenization on a run that qualifies. world_size_from_env coerces anything unusable to 1, so such a run now qualifies again. This is the direction that fails loudest on revert.

Exhaustively: auto, "", 3.5, 1e3, 0x8, None, true all used to veto and now enable. " 8 ", "8\n", "+8", "08" read as 8. 0 and -2 read as 1 under both expressions.

The fix

world_size = world_size_from_env(), with dataset_bounds already imported one line above.

Deliberately not worker.py's _data_parallel_world_size, which also counts visible CUDA devices. The two call sites want different numbers. That one bounds a row subset, where over-counting costs nothing and under-counting recycles rows, so taking the max of env, process group and device count is right. This one feeds a veto, where both directions are wrong, and Studio's own multi-GPU load is device_map="balanced", which transformers treats as model-parallel and pins to _n_gpu = 1: a balanced 4-GPU run draws the same rows per step as one GPU, so counting devices here would report 4x the passes and veto a qualifying run with a fabricated reason. The comment at the call site records this so the two do not get "unified" later.

Also in this PR:

• dataset_bounds.py: world_size_from_rank_files read its bounded prefix from a text handle, where read(n) counts characters, so a rank file of four-byte codepoints could pull four times MAX_WORLD_SIZE_FILE_BYTES off disk. Now read in binary, so the constant means what its name says; json.loads takes bytes and still raises a ValueError on anything that is not UTF-8.
• dataset_bounds.py: deletes a duplicated line that had garbled the world_size_from_rank_files docstring mid-sentence.
• unsloth/utils/packing.py: two comments justified leaving boundary labels unmasked by saying unsloth_zoo "already subtracts the N-1 boundary targets". They now describe the idempotent contract instead, which is what the paired unsloth-zoo change establishes. The code is unchanged and was correct either way.

I checked whether any other site needs the same change. Inside studio/backend/core/training this was the only narrower reader; the other two are world_size_from_env itself and _data_parallel_world_size, and the latter is intentionally wider. resolved_max_steps_epochs has exactly one producer and one consumer in the tree. Nothing else computes an epoch count from max_steps * batch * accum * world_size.

The risk this introduces, quantified

The pass count now depends on eight variables rather than one, so a size variable that is set but not truthful raises resolved_epochs and can veto a run that used to qualify. That is a behaviour change for an existing user on unchanged hardware, and it deserves a number rather than a shrug.

When it can trigger. With S = max_steps * batch * accum and R the rows the trainer sees, the old verdict is S/R <= 1 and the new one S*W/R > 1, so the window is S <= R < S*W. The online-tokenization floor adds R >= 10_000. And the merged row bound, which runs earlier in the same worker process and reads the same variables, cuts any larger corpus to exactly 4*S*W rows, at which point the new pass count is a constant 0.25 and the verdict cannot change at all. So the whole window is a corpus of between max(10_000, S) and S*W rows. With Studio's defaults and a 200-step run under an eight-rank variable, that is 10_000 to 12_799 rows.

What it costs when it triggers. The run takes the eager Arrow map, which is the path that shipped before #8960. No correctness, loss, gradient or checkpoint difference. The opposite direction, under-counting, engages a lazy view that re-tokenizes on every extra pass, so the over-count direction is the safer of the two.

How realistic a stale variable is. I could find no evidence of a login dotfile, module file or mainstream HPC container image exporting these unconditionally. What is real and structural: Slurm's pmi2 and cray_shasta plugins put PMI_SIZE in every task environment and srun defaults to --export=ALL, so anything started from an interactive srun shell inherits it with no MPI in sight. mpirun -np 1 anything sets OMPI_COMM_WORLD_SIZE=1, which is harmless because 1 is the answer either way.

What I did about it. The number is left alone. Narrowing it here, for instance by requiring the RANK partner as routes/inference.py does, would put this call site back into disagreement with the row bound that already reads exactly these variables and already decides which rows the run trains on, and that disagreement is the thing this PR exists to remove. What was actually missing is that the verdict was silent, so the trainer now logs the world size and the variables behind it whenever the environment raises it above 1:

Launcher environment reports 8 data-parallel processes (OMPI_COMM_WORLD_SIZE=8); a step-capped run consumes that many times the rows per step

Nothing is logged on a single-process run.

Launcher contracts, checked against upstream

Several of the comments in dataset_bounds were wrong or imprecise and are corrected here:

PMIX_SIZE and MPI_WORLD_SIZE stay in the tuple, because routes/inference.py already lists them and the two must not disagree about what counts as a launcher, but the comments no longer claim anything sets them. The MLX hostfile test fixture now writes the real shape.

Evidence 1: the decision log line

The observable consequence is a log line, and the trigger cannot be produced from the UI, so there is no screenshot or recording here on purpose. Instead core.training.worker.run_training_process was driven directly, four arms, Qwen3-0.6B LoRA 4bit on unsloth/OpenMathReasoning cot (192,523 rows), max_steps=5000, batch 2, accum 4. The run is killed the moment the line is emitted, so no optimizer step is taken. OMPI_COMM_WORLD_SIZE=8 is injected into the worker's environment before the backend is imported, exactly as mpirun would have left it.

Row two is the bug: an eight-rank mpirun launch drawing 320,000 rows over a 192,523-row split, reported as a single pass. Row three is the no-regression control on the same steps and the same split.

Evidence 2: an old install upgrading, on a single-process run

The proof that matters most for an existing user is that nothing moves when no launcher variable is set. Merge base 27326c4d8 checked out in its own worktree, versus this branch, same config as above, stopped after three optimizer steps:

Identical gate decision, identical decision line, identical loss. grad_norm differs in the fifth decimal, and the third row is why that is not a finding: two runs of the same code differ by more than the two arms do, which is ordinary nondeterministic reduction order on the GPU. Nothing on this branch can touch the arithmetic anyway.

Tests

studio/backend/tests/test_online_tokenization_wiring.py gains an autouse fixture that clears every launcher variable, using the same constant tuples and the same _single_process_launch helper as test_training_preflight.py. That also retires this file's existing dependence on the ambient environment: the pass count is read out of the environment, so before this a case's result depended on what the runner's shell happened to export. Proved load-bearing by exporting WORLD_SIZE=8 OMPI_COMM_WORLD_SIZE=4 and flipping the fixture to autouse = False: 1 failed, 30 passed.

New cases:

1. mpirun (OMPI_COMM_WORLD_SIZE=8) vetoes
2. torchrun's LOCAL_WORLD_SIZE=8 alone vetoes
3. MLX ring hostfile written to tmp_path, in the real per-rank-list shape, vetoes
4. the same payload inline in the variable, in the {"hosts": [...]} object form, vetoes
5. junk WORLD_SIZE (auto, empty, eight) now enables, one parametrised case each
6. a direct assertion on the resolved_max_steps_epochs number handed to the gate, not just its verdict
7. a single-process control at the same step count, which must still qualify
8. the launcher report names the variable, and says nothing at all on a single-process run

Plus, in test_training_preflight.py, the byte cap on a rank-file read (a file with fewer characters than the cap but more bytes than it, which a text handle would have read whole) and the launcher report's own contract, including that it never raises on a hostile mapping.

0 and -4 were deliberately left out of case 5: the old max(1, int(...)) already answered 1 for those, so they would pass against the bug and belong with dataset_bounds' own coercion tests.

Every case was proven discriminating by restoring the old expression and confirming it fails, in the default environment and in two others (python 3.10 / torch 2.6.0 / transformers 4.57.6 / TRL 0.22.2, and python 3.13 / torch 2.13.0 / transformers 5.15.0 / TRL 1.10.0), with an identical failure set each time:

$ git apply -R <the trainer.py hunk>
$ PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_online_tokenization_wiring.py -q
FAILED test_an_mpirun_launch_scales_the_rows_a_step_consumes
FAILED test_a_per_node_torchrun_scales_the_rows_a_step_consumes
FAILED test_an_mlx_hostfile_scales_the_rows_a_step_consumes
FAILED test_an_inline_hosts_payload_scales_the_rows_a_step_consumes
FAILED test_a_junk_world_size_no_longer_disables_online_tokenization[auto]
FAILED test_a_junk_world_size_no_longer_disables_online_tokenization[]
FAILED test_a_junk_world_size_no_longer_disables_online_tokenization[eight]
FAILED test_the_resolved_pass_count_handed_to_the_gate_is_the_arithmetic
8 failed, 23 passed

test_a_single_process_launch_still_qualifies correctly stays green there: it is the control, world size 1 either way. Reverting the byte-cap hunk fails test_a_rank_file_read_is_capped_in_bytes_not_characters, and removing the launcher report fails its own case.

With the fix in place, and unchanged with CUDA_VISIBLE_DEVICES="":

test_online_tokenization_wiring.py + test_training_preflight.py + test_online_tokenization.py
176 passed

File order was permuted three ways and the same three files were run under pytest-randomly seeds 11, 22 and 33 in two separate interpreters: 172 passed every time at the point that was measured, so the new autouse fixture leaks nothing into its neighbours and nothing leaks into it.

A pre-existing test defect this surfaced

The first cross-platform staging run failed on macos-14 and windows-latest with not enough CPU workers to stay ahead of the GPU. resolve_worker_count sizes itself from CPU affinity and the cgroup quota and returns 0 when the host cannot spare MIN_ONLINE_WORKERS, and that gate sits ahead of the pass-count veto, so on a small runner every case in the file was asserting about the runner rather than about the wiring.

Reproduced locally with taskset -c 0, which is the same condition:

So the failure predates this branch; it is simply that this branch adds cases to a file that was already host-dependent. _run now pins resolve_worker_count for the same reason it already pins the TRL feature detector. The worker gate keeps its own coverage in test_online_tokenization.py.

Compatibility, and what was actually executed

dataset_bounds.py stays torch-free, as its module docstring requires; asserted mechanically by importing it in a bare interpreter and checking that none of torch, transformers, trl, datasets, unsloth, unsloth_zoo or numpy enter sys.modules. The changed code uses only os, json and builtins, so no library version can change its answer; the version exposure is the test file, which imports the real trainer.

Eight stack combinations were executed on Linux, covering python 3.10 / 3.11 / 3.12 / 3.13, transformers 4.57.6 and 5.15.0, TRL 0.22.2 / 0.25.1 / 1.10.0, and torch 2.6.0 / 2.9.1 / 2.13.0, arranged so that all four (transformers, TRL) combinations appear and python and torch are crossed rather than held constant. Every cell reported a real pass count with no module-level skip and no importorskip firing, checked by probing that datasets, torch and the real trl all imported and that the live trl_supports_skip_prepare_dataset() returned True on every TRL. Note that transformers 5.x and torch 2.13 sit outside the range studio-backend-ci.yml declares, so those cells are future-proofing rather than evidence about the CI job.

The os.path.isfile on an arbitrary environment string was exercised against a directory, a fifo, a symlink loop, a dangling symlink, a mode-000 file as a non-root user, a 2 GB file, invalid UTF-8, /dev/zero, /dev/null, a missing path and a path containing spaces and non-ASCII characters. All return 1 without raising, the slowest case in the whole matrix is 2.7 ms, and nothing blocked. Payloads covered the bare list, the {"hosts": [...]} object, an empty list, a JSON string, a JSON number, null, a truncated payload, an inline payload above the cap and two rank-file variables disagreeing.

Honest limits. Windows, macOS and WSL were exercised as sys.platform fakes and through cross-platform staging CI, not as native development runs; on both of those platforms the online-tokenization platform gate vetoes before the pass count is read at all, so the changed expression cannot alter a decision there. A UNC path on Windows and a hung NFS mount on Linux are the one residual risk I could not clear: the block would be in the stat, which nothing here bounds, and I will not stage a fake hang and call it evidence. AMD/ROCm was not executed; the decision path reads no device state, which is an argument rather than a run. No real torchrun, mpirun or mlx.launch was involved; every world size in the tests is an injected variable or a temp-file payload. No browser is involved anywhere in this diff, so there is no browser testing to report.

Cross-platform staging CI, test_online_tokenization_wiring.py + test_training_preflight.py with torch>=2.4 datasets trl>=0.24 installed so the tests cannot importorskip themselves into a vacuous exit 5: