docs/TROUBLESHOOTING.md

# TROUBLESHOOTING — Wave 14

This document catalogs every Wave-14-known failure mode in the Composer
Replication Framework, along with how to diagnose, fix, and verify each
one. It is intentionally surgical: the surface area added in Waves 12–14
(SimPO/TAID/Entropy-OPD distillation kwargs, the PRIME-RL composer-loss
adapter, the serverless DiLoCo `MockManager` + `ObjectStoreAllReduce`
path, and the data-juicer-backed replaysim normalizer) introduced new
ways for users to trip themselves up. Each failure mode here is something
a maintainer has actually seen or anticipated during the cross-model
review of Wave 14.

If you hit something not covered below, jump to the
[How to file a bug report](#how-to-file-a-bug-report) section at the end —
the template there gives a maintainer everything they need to reproduce.

---

## Common things to check first

Before reading any further, run through this checklist. ~80% of "framework
broken" reports turn out to be one of these:

1. **Python version.** The framework targets Python 3.10–3.12. The
   `pyproject.toml` `target-version` is `py310`. If you are on 3.13+,
   transitive deps (notably Ray, pulled in by data-juicer) may not yet
   ship wheels and will try to build from source. Run `python --version`.

2. **Fresh virtual environment.** Mixing the framework into an existing
   environment that already has `torch`, `transformers`, `trl`, or
   `torchft` pinned to incompatible versions is the #1 source of import-
   time errors. Create a new venv: `python -m venv .venv && source
   .venv/bin/activate && pip install -e .[dev]`.

3. **Editable install.** Most contributors run `pip install -e .` so
   that local edits to `composer_replication/` are picked up. If you
   `pip install composer-replication` from a registry instead, your
   edits to the source tree will be ignored. Confirm with
   `pip show composer-replication | grep Location`.

4. **Optional extras.** Several modules are optional-dep gated:
   - `[replay]` — adds `httpx` (used for OpenRouter teacher calls).
   - `[train]` — adds TRL, peft, accelerate, datasets (production GRPO).
   - `[replaysim]` — adds `data-juicer` (and via it, Ray as a transitive).
   - `[serverless]` — adds `fsspec`. For non-local rendezvous URIs you
     also need a backend-specific fsspec adapter (see Failure Mode 5).
   - `[dev]` — adds `pytest`, `ruff`, etc.
   If you see `ModuleNotFoundError: No module named 'data_juicer'`, you
   forgot the extra. Install with `pip install -e .[replaysim]`.

5. **Run the test suite first.** Before debugging anything, run the
   subset of tests touching the area you care about:
   ```
   pytest composer_replication/tests/                 # core compose_loss
   pytest composer_replication/distillation/tests/    # SimPO / TAID / OPD
   pytest composer_replication/recipes/prime_rl/tests/  # PRIME-RL adapter
   pytest composer_replication/diloco/serverless/tests/ # MockManager + DiLoCo
   pytest composer_replication/replaysim/tests/       # data-juicer normalizer
   ```
   If any green test fails for you locally, the problem is environmental
   — fix that before digging into your own code.

6. **Read the docstring of the symbol you're calling.** Wave 14
   docstrings are written to be the first line of documentation. The
   `compose_loss` docstring (`composer_replication/loss.py`) lists every
   required and optional input key. The `MockManager` docstring
   enumerates the torchft surface methods it implements.

---

## Failure modes

### 1. `pip install -e .[replaysim]` hangs or fails on Python 3.12 with a Ray-related path error

**SYMPTOM.** Installing the `[replaysim]` extra (which pulls
`data-juicer`) triggers a transitive install of Ray. On Python 3.12, the
first `import ray` (often during `pip` build hooks or the first time
data-juicer is loaded) fails with messages mentioning
`/tmp/ray/session_*` paths, missing `pyarrow` symbols, or `OSError:
[Errno 2] No such file or directory: '/dev/shm/ray-...'` inside Docker.

**DIAGNOSIS.** `data-juicer` declares `ray` as a transitive dependency.
On Python 3.12 the wheel matrix is incomplete for some Ray versions, and
Ray's first-import probes `/dev/shm` and `/tmp/ray` for its session
state. In a sandboxed container, restricted CI runner, or WSL
environment with a non-default `/tmp`, those probes fail. Wave 14
subagent T2 hit this in CI and worked around it by pinning Ray and by
making sure `/tmp` exists and is writable.

**FIX.**
- Prefer Python 3.11 if you're on 3.12+ and don't need 3.12 features.
- If you must stay on 3.12, ensure `/tmp` is writable and pre-create the
  session directory: `mkdir -p /tmp/ray && chmod 1777 /tmp/ray`.
- In Docker, mount a real tmpfs at `/dev/shm`:
  `docker run --shm-size=2g …`.
- If you don't need replaysim normalization, you can skip the extra
  entirely. The `DJNormalizer(skip_dj=True)` passthrough (see
  `composer_replication/replaysim/normalize.py:165`) does not import
  `data_juicer` and therefore does not import Ray.

**VERIFICATION.** The skip-dj passthrough is exercised by
`test_dj_normalizer_skip_dj_passthrough` and
`test_dj_normalizer_skip_dj_preserves_count` in
`composer_replication/replaysim/tests/test_replaysim.py`. Both run
without `data_juicer` installed:

```
pytest composer_replication/replaysim/tests/test_replaysim.py::test_dj_normalizer_skip_dj_passthrough -xvs
```

If that passes in your environment, your `[replaysim]`-less install is
healthy — only the full data-juicer code path requires Ray.

---

### 2. `compose_loss` produces wrong-looking numbers when combining new kwargs

**SYMPTOM.** You pass several Wave-14 distillation kwargs to
`compose_loss` (e.g. `dpo_variant="simpo"`, `sdpo_wrapper="taid"`,
`taid_schedule_step=0`, `simpo_beta=2.0`, `entropy_opd_h_max=…`), and
the loss curve looks wrong: NaNs, identically-zero `sdpo_jsd` channel,
or a `total` that is bit-different from your reference run with no
distillation kwargs at all.

**DIAGNOSIS.** `compose_loss` now has 13 keyword arguments and the
contract between them is non-trivial. Subagent T1's review identified
three combinations that look reasonable but are unsupported:
- Passing `taid_schedule_step` without `taid_total_steps` (or vice
  versa). The function raises `ValueError` clearly, but the message can
  scroll past in noisy logs.
- Passing `dpo_variant="simpo"` while still supplying
  `dpo_chosen_ref_logprobs`. Those keys are **silently ignored** —
  SimPO is reference-free.
- Passing `sdpo_wrapper="taid"` without supplying either
  `student_init_logits` OR `student_init_input_ids` in `inputs`. The
  function will fall back to a forward pass through the (possibly
  drifted) live model, which is a footgun late in training (see Failure
  Mode 8).

**FIX.** Read the docstring at the top of
`composer_replication/loss.py` (lines 25–39 list the three pluggable
losses and their preconditions). The general rule:

```python
from composer_replication import compose_loss

# Defaults (no distillation knobs) reproduce legacy 3-channel composition bit-exact.
out = compose_loss(model, inputs)

# To opt into SimPO, pass dpo_variant ONLY. Do not pass ref-logprob keys.
out = compose_loss(model, inputs, dpo_variant="simpo",
                   simpo_beta=2.0, simpo_gamma=1.0)

# To opt into TAID, pass BOTH schedule_step AND total_steps, AND make sure
# inputs["student_init_logits"] is populated (see Failure Mode 8).
out = compose_loss(model, inputs, sdpo_wrapper="taid",
                   taid_schedule_step=step, taid_total_steps=total_steps)
```

Setting all 13 kwargs to their defaults is **bit-exact equivalent** to
the pre-Wave-13 3-channel loss; if your defaults call gives different
numbers than your old code, file a bug.

**VERIFICATION.** The bit-exact equivalence and every supported
combination is locked in by the 11 integration tests in
`composer_replication/tests/test_compose_loss_integration.py`. The most
important ones:
- `test_defaults_bit_exact_with_legacy_kwargs` — passing the new kwargs
  at their defaults is identical to legacy.
- `test_simpo_does_not_require_ref_logprobs` — SimPO works with the
  ref-logprob keys absent from `inputs`.
- `test_taid_alpha_one_recovers_sdpo` — TAID with `alpha_min=alpha_max=1`
  reproduces standard SDPO.
- `test_taid_requires_schedule_step` / `test_taid_requires_total_steps` —
  the partial-config error path.

```
pytest composer_replication/tests/test_compose_loss_integration.py -xvs
```

---

### 3. `MockManager` works today but silently breaks after a torchft upgrade

**SYMPTOM.** Your serverless DiLoCo run starts, the first outer round
completes, and then `torchft.DiLoCo` raises an `AttributeError` on
something like `_use_async_quorum`, `should_commit`, or
`current_step` — or worse, it silently uses the wrong sync semantics.

**DIAGNOSIS.** `MockManager` is a duck-typed shim that mirrors
`torchft.Manager` rather than subclassing it. The surface it implements
is enumerated in the docstring at
`composer_replication/diloco/serverless/allreduce.py:215`:

> Methods/attributes DiLoCo touches: `allreduce`, `should_commit`,
> `start_quorum`, `current_step`, `disallow_state_dict_read`,
> `allow_state_dict_read`, `register_state_dict_fn`, `_use_async_quorum`
> (attribute), `num_participants`, `rank`.

The two **private** members in that list — `_use_async_quorum` and the
internal `current_step` counter — are private torchft API that may be
renamed without notice in any torchft minor release. Wave 14 subagent
T3 specifically called this out: "If torchft renames `_use_async_quorum`
to anything else, MockManager silently breaks because there is nothing
holding the contract beyond a string."

**FIX.**
- **Pin torchft.** In `pyproject.toml` keep your torchft version pinned
  to a known-good range (e.g. `torchft>=0.2,<0.4`). When you need to
  upgrade, do so deliberately and re-run the integration tests below
  before merging.
- **Watch the deprecation warning.** Wave 14 sets up a clear path to
  warn if `_use_async_quorum` is read on a fresh instance — see the
  comment at `allreduce.py:255`.
- **Don't pass an arbitrary torchft branch.** If you've patched torchft
  locally, the `MockManager` may need updating in lockstep. The
  surface-compatibility tests below will catch this in CI.

**VERIFICATION.** The full DiLoCo × MockManager surface is exercised by:
- `test_mock_manager_shape_compat` in
  `composer_replication/diloco/serverless/tests/test_serverless_local.py`
  — sanity check that all expected methods/attributes exist.
- `test_mockmanager_has_full_diloco_call_surface` in
  `composer_replication/diloco/serverless/tests/test_serverless_diloco_integration.py`
  — runs an end-to-end outer round through real torchft `DiLoCo`,
  hitting every method on the surface list above.
- `test_mockmanager_diloco_outer_round_completes` — full one-round
  smoke ending in a successful outer SGD step.

If any of these tests turn red after a torchft bump, **do not ship**:
inspect the new torchft Manager surface and update `MockManager`
to match.

```
pytest composer_replication/diloco/serverless/tests/test_serverless_diloco_integration.py -xvs
```

---

### 4. SimPO loss curve looks like noise

**SYMPTOM.** You wired in `dpo_variant="simpo"`, the run starts, and
the `trace_replay_dpo` channel either drifts to large negative values
(→ `total` blows up) or oscillates with much higher variance than
standard DPO. The loss curve "looks like noise."

**DIAGNOSIS.** SimPO uses **average per-token log-probability**
(`Σ logπ(c_t) / |c|`), not sum log-prob. From the SimPO docstring
(`composer_replication/distillation/simpo.py:11–18`):

> SimPO drops the reference-policy term, replaces it with a target
> margin γ, and uses **average sequence log-probability instead of
> sum**. […] L_SimPO = -log σ( β · [avg_logπ(c) - avg_logπ(r)] - γ )

If you compute `chosen_logprobs.sum()` (or any unmasked aggregation) and
hand it to SimPO as `chosen_avg_logprobs`, the loss is undefined: β=2.0
times a sum-log-prob is on a totally different scale than β=2.0 times an
average. The result looks plausible per-batch but the optimum is
nowhere near the dataset's true preference signal.

**FIX.** Use the helper
`composer_replication.distillation.simpo.avg_sequence_logprob`:

```python
from composer_replication.distillation.simpo import (
    simpo_loss, avg_sequence_logprob,
)

chosen_avg = avg_sequence_logprob(chosen_logprobs, chosen_response_mask)
rejected_avg = avg_sequence_logprob(rejected_logprobs, rejected_response_mask)
loss = simpo_loss(chosen_avg, rejected_avg, beta=2.0, gamma=1.0)
```

The mask is **1 on response tokens, 0 on prompt+padding** — same
convention as the rest of the framework. If you must roll your own
aggregation, divide by `response_mask.sum(dim=-1).clamp_min(1.0)`,
not by `response_mask.shape[-1]`.

**VERIFICATION.** The avg-vs-sum semantics are pinned by
`test_avg_sequence_logprob` in
`composer_replication/distillation/tests/test_distillation_losses.py`,
which constructs known per-token log-probs and asserts the helper
returns the correct per-sequence average. The end-to-end SimPO
loss-shape check is `test_simpo_loss_returns_scalar` in the same file.

```
pytest composer_replication/distillation/tests/test_distillation_losses.py::test_avg_sequence_logprob -xvs
pytest composer_replication/distillation/tests/test_distillation_losses.py::test_simpo_loss_lower_for_better_separation -xvs
```

---

### 5. `ObjectStoreAllReduce` works locally but fails on `s3://` at first allreduce

**SYMPTOM.** You construct
`ObjectStoreAllReduce(uri="s3://my-bucket/run42/", rank=0,
world_size=4)`. The constructor succeeds. The first call to
`allreduce(tensor, name="...")` raises `ImportError: Install s3fs to
access S3` or `botocore.exceptions.NoCredentialsError: Unable to locate
credentials`.

**DIAGNOSIS.** `ObjectStoreAllReduce` uses fsspec to reach the
backend, but **fsspec only ships protocol stubs, not adapters**. The
constructor doesn't know which protocol you'll use and doesn't
eagerly validate, so it accepts any URI. The `s3://` adapter requires:
1. The `s3fs` package (`pip install s3fs`), which is **not** in the
   default `[serverless]` extra.
2. Working AWS credentials (env vars, `~/.aws/credentials`, IAM role,
   or whatever your environment normally provides to boto3).

The same is true for `gs://` (`gcsfs`), `az://` (`adlfs`), and
`hf://` (`huggingface_hub`'s fsspec integration, which is included if
you have `huggingface_hub` installed).

**FIX.**
- Install the right adapter alongside the framework:
  ```
  pip install s3fs       # for s3://
  pip install gcsfs      # for gs://
  pip install adlfs      # for az://
  ```
- Verify credentials work outside the framework first:
  ```
  python -c "import s3fs; print(s3fs.S3FileSystem().ls('my-bucket'))"
  ```
- If you're running on Modal/HF Jobs, set the credentials as Modal
  secrets / HF Jobs env vars in the executor config — not in your
  local shell.

The constructor could in principle perform an eager probe (e.g. a
`HEAD` on the rendezvous prefix) to fail fast at init time. Wave 14
deliberately did not add this because it adds a network round-trip on
every replica startup. If you want pre-flight validation in your
training script, call `fsspec.filesystem(protocol).ls(uri)` yourself
before constructing the manager.

**VERIFICATION.** The `file://` and bare-path code paths — the only
ones that don't need an extra adapter — are exercised by:
- `test_object_store_allreduce_local_paths_create_dir`
- `test_object_store_allreduce_world_size_1_passthrough`
- `test_object_store_allreduce_round_id_increments`

…all in
`composer_replication/diloco/serverless/tests/test_serverless_local.py`.
If those pass and your `s3://` URI fails, the framework is fine and
your fsspec adapter or credentials are the problem.

```
pytest composer_replication/diloco/serverless/tests/test_serverless_local.py -xvs
```

---

### 6. Custom replaysim recipe drops every record (or crashes data-juicer)

**SYMPTOM.** You wrote a custom replaysim YAML recipe modeled on
`composer_replication/recipes/replaysim/default.yaml`. It loads
without error, but every input DPO pair is dropped, OR data-juicer
raises `KeyError: 'text_key'`, OR it raises a complaint about
"expected str, got list" inside one of the filters.

**DIAGNOSIS.** Wave 14 fixed two related bugs in the *default* recipe
that custom-recipe authors will hit again. Both are documented in the
header comment at
`composer_replication/recipes/replaysim/default.yaml:21–35`:

1. **`text_keys` plural vs `text_key` singular.** The top-level
   dataset contract uses `text_keys: chosen` (plural). Each individual
   op uses `text_key: chosen` (singular). They are not interchangeable.
   data-juicer's dataset loader validates that the `text_keys` field
   exists on every record before any op runs; an op that uses
   `text_keys` instead of `text_key` is silently misconfigured.

2. **`chosen` / `rejected` as strings vs as list-of-dicts.**
   data-juicer ops like `text_length_filter`, `words_num_filter`,
   `special_characters_filter`, and `document_deduplicator` read a
   single string field. Pointing them at the chat-messages list
   (`chosen_messages`, `rejected_messages`) crashes or silently
   no-ops. The framework's `_dpo_pair_to_dj_record` keeps **both**
   shapes side-by-side: `chosen`/`rejected` (strings) for filter ops,
   and `chosen_messages`/`rejected_messages` (chat-messages list) for
   chat-aware ops + the `NormalizedDPOPair` round-trip.

**FIX.** Treat the default recipe as your starting template. Concretely:
- Always declare `text_keys: chosen` at the top.
- For every length/word/special-char op you add, duplicate it: once
  with `text_key: chosen`, once with `text_key: rejected`. (Each op
  takes only one `text_key` — see comment at lines 31–35 of
  `default.yaml`.)
- Never point a filter op at `chosen_messages` or `rejected_messages`.
  Those are list-of-dicts; only chat-aware ops accept that shape.

**VERIFICATION.** The two-shape contract is locked in by:
- `test_record_chosen_rejected_are_flat_strings_for_dj_text_ops` —
  asserts `chosen` and `rejected` are bare strings on every record
  produced by `_dpo_pair_to_dj_record`.
- `test_record_chosen_rejected_messages_carry_chat_shape` — asserts
  `chosen_messages` / `rejected_messages` exist as list-of-dicts.
- `test_dj_normalizer_e2e_default_recipe(tmp_path)` — runs the actual
  default recipe through real data-juicer end-to-end (skipped if
  `data_juicer` isn't importable).

…all in
`composer_replication/replaysim/tests/test_replaysim.py`. If those
pass and your custom recipe still drops everything, diff your YAML
against `default.yaml` until the two shapes align.

```
pytest composer_replication/replaysim/tests/test_replaysim.py -xvs
```

---

### 7. `ValueError: expected (seq,) shape, got (B, T)` from PRIME-RL composer_loss

**SYMPTOM.** You wired the PRIME-RL recipe into a training loop you
adapted from another framework (TRL, openrlhf, etc.), and on the very
first `loss_fn` call you get a `ValueError` mentioning shape
`(seq,)` versus `(B, T)`.

**DIAGNOSIS.** PRIME-RL calls its loss function **one sample at a
time**, with 1-D `(seq,)` tensors — not batched `(B, T)` tensors. The
recipe's docstring spells this out at
`composer_replication/recipes/prime_rl/composer_loss.py:16–30`:

> Note the **per-sample (seq,) shape** — PRIME-RL's runner calls the
> loss function one sample at a time, not on a batched (B, T) tensor.

Wave 14 fixed an earlier draft of the recipe that incorrectly assumed
`(B, T)`. The new version raises a clear `ValueError` if you hand it
the wrong shape, instead of silently broadcasting and producing
nonsense gradients. Users who are used to TRL or openrlhf — both of
which call the loss with batched tensors — see this on day one.

**FIX.**
- If you are running inside PRIME-RL via its `CustomLossConfig`, you
  don't need to do anything: PRIME-RL's runner produces `(seq,)`
  tensors and the recipe accepts them.
- If you are calling the recipe directly from your own runner, slice
  your batch into per-sample 1-D tensors before each call:
  ```python
  for b in range(B):
      inputs_b = LossInputs(
          trainer_logprobs=batched.trainer_logprobs[b],
          inference_logprobs=batched.inference_logprobs[b],
          advantages=batched.advantages[b],
          loss_mask=batched.loss_mask[b],
          teacher_logprobs=None if batched.teacher_logprobs is None
                          else batched.teacher_logprobs[b],
      )
      loss = loss_fn(inputs_b, ...)
  ```
- If you genuinely need a batched API, write a thin wrapper around
  `loss_fn`. Don't patch the recipe — its shape contract is dictated
  by PRIME-RL, not by us.

**VERIFICATION.** The shape contract is pinned by two tests in
`composer_replication/recipes/prime_rl/tests/test_composer_loss.py`:
- `test_advantages_shape_validates_seq_accepted` — `(seq,)` succeeds.
- `test_advantages_shape_validates_bt_rejected` — `(B, T)` raises
  `ValueError`.

```
pytest composer_replication/recipes/prime_rl/tests/test_composer_loss.py -xvs
```

---

### 8. TAID can't run mid-training because `student_init_logits` is missing

**SYMPTOM.** You decide partway through a training run to enable
`sdpo_wrapper="taid"` (e.g. you read the TAID paper after step 2000
and want to retrofit). The next training step blows up — either with
a `KeyError` for `student_init_logits` / `student_init_input_ids`, or
with a strange-looking loss because the framework fell back to
re-running a forward pass through the *current* (drifted) model
instead of the init model.

**DIAGNOSIS.** TAID interpolates between the **student's distribution
at step 0** and the teacher's distribution. From the TAID docstring at
`composer_replication/distillation/taid.py:10–24`:

> TAID interpolates between an "identity" target (the student's own
> distribution at step 0) and the teacher's distribution, with the
> interpolation coefficient annealed from 0 → 1 over training.

That step-0 reference target has to come from somewhere. The framework
accepts it via either:
1. `inputs["student_init_logits"]` — a precomputed `(B, T, V)` tensor
   captured at training start (preferred for production), OR
2. `inputs["student_init_input_ids"]` — input ids for a frozen forward
   pass through `model`. **This assumes `model` has not yet drifted
   from init.** It is correct only at step 0 or in tests; in
   production it silently produces the wrong target.

If you forgot to capture the init logits at step 0, you cannot
faithfully use TAID mid-run.

**FIX.** Capture init logits at step 0 and persist them:

```python
# At step 0, before any optimizer.step() call:
with torch.no_grad():
    init_logits = model(input_ids=batch["input_ids"]).logits
    # Save to disk if you'll need them across restarts:
    torch.save(init_logits, "checkpoints/init_logits_batch0.pt")
    inputs["student_init_logits"] = init_logits

# Or, if you have a fixed eval probe set, capture init logits once
# for that fixed set and reuse them every step:
inputs["student_init_logits"] = cached_init_logits
```

If you genuinely have no step-0 snapshot, **TAID is not retrofittable**
to your run. Your options are:
- Restart from a checkpoint that *was* the step-0 model.
- Use a different distillation wrapper (`sdpo_wrapper="entropy_opd"`)
  that doesn't need init logits.
- Accept the bias from the live-model fallback path. Don't.

**VERIFICATION.** The precomputed-vs-live-fallback contract is exercised by:
- `test_taid_accepts_precomputed_student_init_logits` in
  `composer_replication/tests/test_compose_loss_integration.py` —
  passes precomputed logits and asserts the TAID-wrapped channel uses
  them.
- `test_taid_alpha_one_recovers_sdpo` — asserts that with
  `alpha_min=alpha_max=1.0` (i.e. pure teacher target, init logits
  ignored) TAID reproduces standard SDPO. If your training ignores
  init logits silently, *this* is the test that would have failed.

```
pytest composer_replication/tests/test_compose_loss_integration.py::test_taid_accepts_precomputed_student_init_logits -xvs
```

---

### 9. `ModalExecutor()` or `HFJobsExecutor()` raises `NotImplementedError` at construction

**SYMPTOM.** You write
`executor = ModalExecutor(app_name="my-app")` (or the HF Jobs
equivalent) in a production script and the constructor immediately
raises:

```
NotImplementedError: ModalExecutor is a v0 skeleton; full implementation pending.
Use LocalProcessExecutor for testing.
```

Same for `HFJobsExecutor`. This is at *init time*, not at the first
`launch_replicas` call.

**DIAGNOSIS.** Per ADR-005 the v0 release ships only the
`ServerlessExecutor` Protocol and the reference `LocalProcessExecutor`.
The Modal and HF Jobs implementations are **import-safe skeletons** —
the classes exist and you can `from … import ModalExecutor`, but
`__init__` raises `NotImplementedError` to prevent silent partial
behavior. See `modal.py:64` and `hf_jobs.py:64`.

This is intentional. We didn't want to ship a half-working Modal
executor that succeeds at `launch_replicas` and then silently fails
two-thirds of the way through `collect`.

**FIX.**
- Use `LocalProcessExecutor` for development, CI, and any single-host
  multi-process testing.
- For real cloud deployment in the v0 era, run your training script
  directly in Modal/HF Jobs by hand: write your own thin Modal
  function that constructs `MockManager(ObjectStoreAllReduce(uri,
  rank, world_size))` and runs the training loop. The skeleton
  docstrings at `modal.py:24–48` and `hf_jobs.py:26–49` show exactly
  the pattern.
- Watch the `BACKLOG.md` for v0 polish — the real implementations are
  scheduled.

**VERIFICATION.** That `LocalProcessExecutor` is fully functional and
correctly implements the Protocol is locked in by:
- `test_local_executor_runs_allreduce_across_replicas` in
  `composer_replication/diloco/serverless/tests/test_serverless_local.py`
  — runs N replicas locally, performs an allreduce across them.
- `test_local_executor_handles_multiple_rounds`
- `test_local_executor_reports_failed_replicas`

If those tests pass, your serverless DiLoCo machinery works — only the
specific cloud adapters are missing. The skeletons themselves are not
under test (raising in `__init__` is the contract).

```
pytest composer_replication/diloco/serverless/tests/test_serverless_local.py -xvs
```

---

### 10. DPPO mask drops every token — "loss became 0" or "no gradients"

**SYMPTOM.** You ported a PPO config from another framework (KL
penalty + clip ε=0.2 + value loss), wired it into the PRIME-RL recipe
with the default `dppo_mask_high=0.2` / `dppo_mask_low=0.2`, and the
training loss is suspiciously close to zero. Inspecting the recipe's
internal `keep_mask` shows nearly every token is being masked out.

**DIAGNOSIS.** PRIME-RL's "DPPO mask" is **not** the same as PPO
clipping, and not even the same as a log-ratio threshold. From the
recipe docstring at
`composer_replication/recipes/prime_rl/composer_loss.py` (mirroring
PRIME-RL upstream `prime_rl/trainer/rl/loss.py` lines 137-148):

> The mask gate is on **probability-space**
> `probs_diff = exp(trainer_lp) - exp(inference_lp)`, NOT on the
> log-ratio. A positive-advantage token is dropped iff
> `probs_diff > dppo_mask_high`; a negative-advantage token iff
> `probs_diff < -dppo_mask_low`. Masked tokens are **dropped from the
> policy-gradient term** but still contribute to the KL penalty.

The defaults `dppo_mask_high=dppo_mask_low=0.2` match PRIME-RL's
`DefaultLossConfig`. Because the gate is on probability-space, the
"in-band" zone is
`exp(trainer_lp) ∈ [exp(inference_lp) - 0.2, exp(inference_lp) + 0.2]`.
For a token with inference probability ~0.5 this is a fairly tight
band; for tokens at probability ~0.001 or ~0.999 the same threshold
behaves very differently from a log-ratio bound. This is by design —
PRIME-RL is bounding the absolute change in token probability, not the
multiplicative change.

The two failure modes:

1. **All tokens masked.** Trainer and inference engines disagree
   sharply (fp16 vs bf16, stale rollout cache, mismatched chat
   templates) and `probs_diff` exceeds 0.2 almost everywhere.
2. **No tokens masked.** Trainer ≈ inference (e.g. you forgot to step
   the optimizer between rollouts) so the bound is never binding and
   the policy never sees any DPPO regularization.

**FIX.** Inspect the empirical `probs_diff` distribution before
tuning:

```python
# In your training loop:
probs_diff = torch.exp(trainer_logprobs) - torch.exp(inference_logprobs)
print(torch.quantile(probs_diff.abs(), torch.tensor([0.5, 0.9, 0.99])))
```

For a healthy on-policy run with bf16 trainer + bf16 inference and
fresh rollouts, the central 99% of `|probs_diff|` should sit well
below `0.2`. If yours doesn't, the upstream divergence is the
problem, not the bound. Bumping `dppo_mask_high/low` to 0.5 or 1.0 is
a workaround but it disables the trust-region intent of DPPO.

**Do not** translate PPO ε=0.2 directly. PPO ε=0.2 is a multiplicative
log-ratio bound (`|log_ratio| < log(1.2) ≈ 0.18`); DPPO's 0.2 is an
**additive probability-space** bound. The semantics are different and
the defaults are deliberately tight in probability space.

If you genuinely want to disable the mask (e.g. for bug-isolation),
pass `dppo_mask_high=1e6, dppo_mask_low=1e6` (both are
`Field(..., ge=0)` upstream — negative values are rejected by
both PRIME-RL and our adapter). There is a regression test for
exactly this knob.

**VERIFICATION.**
- `test_dppo_mask_high_drops_positive_advantage_outliers` and
  `test_dppo_mask_low_drops_negative_advantage_outliers` in
  `composer_replication/recipes/prime_rl/tests/test_composer_loss.py`
  — assert that out-of-bound tokens are dropped from the
  policy-gradient term (with the upstream sign-of-advantage gate).
- `test_dppo_mask_sign_conditioned_on_advantage` — asserts that a
  positive-advantage token with a large *negative* probs_diff is NOT
  dropped (PRIME-RL only checks the upper bound for positive-advantage
  tokens).
- `test_dppo_bounds_can_be_disabled` — asserts that very wide bounds
  (`1e6`) pass every token through.
- `test_parity_with_prime_rl_default_loss_fn` — when `prime-rl` is
  installed, runs identical inputs through PRIME-RL upstream and our
  adapter and asserts the loss matches.

```
pytest composer_replication/recipes/prime_rl/tests/test_composer_loss.py -xvs
```

---

### 11. `compose_loss` runs but the GRPO channel doesn't behave like real GRPO

**SYMPTOM.** You read the README, saw the "3-channel composition: GRPO
+ SDPO + trace-replay DPO" tagline, called `compose_loss(model,
inputs)` directly in your training loop, and your reward curve never
moves the way it would in a real GRPO trainer. Or: you compared
against a TRL `GRPOTrainer` baseline and `compose_loss` produces
totally different numbers.

**DIAGNOSIS.** From the docstring at the top of
`composer_replication/loss.py:1–16`:

> This is a verification-harness mirror of
> `ComposerReplicationTrainer._compute_loss` that does NOT depend on
> TRL's GRPOTrainer parent. The GRPO channel is replaced with standard
> LM next-token-prediction cross-entropy, which is the limit GRPO
> converges to under deterministic rewards.
>
> Use it for: CPU smokes on real HF models, unit tests of loss
> composition without spinning up TRL, anywhere we want to verify
> gradient flow through the 3-channel sum without paying TRL's full
> machinery cost.
>
> **Do NOT use it as the production training loss.** Production =
> ComposerReplicationTrainer (a real GRPOTrainer subclass).

The `lm_ce` channel labelled "GRPO" in the LossComponents dataclass is
a **stub**: it is plain language-modeling cross-entropy. It is the
correct channel for verification (gradient flow, channel weighting,
distillation wiring), but it is not GRPO's surrogate objective and
will never produce the same numbers as real GRPO under stochastic
rewards.

Real GRPO requires:
- A reward model or rule-based reward,
- Per-prompt advantage estimation across G samples,
- An importance-sampling-ratio clip / mask.

Those live in TRL's `GRPOTrainer`, in our PRIME-RL recipe at
`composer_replication/recipes/prime_rl/composer_loss.py`, or (when
shipped) in a future VeRL recipe.

**FIX.**
- For production GRPO training, do **not** call `compose_loss` directly.
  Instead use one of:
  - `composer_replication.trainer.composer_trainer.ComposerReplicationTrainer`
    — TRL `GRPOTrainer` subclass, full machinery.
  - `composer_replication.recipes.prime_rl.composer_loss.loss_fn` —
    PRIME-RL's `CustomLossConfig` adapter (channel 1 is real DPPO-clipped GRPO).
- For ablations, smokes, and unit tests, `compose_loss` is the right
  tool — but log the `lm_ce` channel as `lm_ce`, not as `grpo`. The
  `LossComponents` dataclass already names the field correctly; if
  your wandb logger relabels it as "GRPO loss", fix the label.

**VERIFICATION.**
- The 11-test integration suite at
  `composer_replication/tests/test_compose_loss_integration.py` only
  asserts gradient flow + bit-exact composition; it deliberately does
  not assert any GRPO-specific property of `compose_loss`. That's the
  contract.
- The PRIME-RL recipe's real DPPO+KL behavior is asserted by
  `test_returns_finite_scalar`,
  `test_dppo_mask_high_drops_positive_advantage_outliers`,
  `test_dppo_mask_sign_conditioned_on_advantage`, and
  `test_parity_with_prime_rl_default_loss_fn` (skip-marked when
  `prime-rl` is not installed)
  in `composer_replication/recipes/prime_rl/tests/test_composer_loss.py`.
  Those tests verify a real importance-sampling-ratio gradient with
  PRIME-RL's advantage-conditioned mask, which `compose_loss` would
  not pass.

If you find yourself wanting `compose_loss` to behave like real GRPO,
that is the signal to switch to one of the production paths above.

```
pytest composer_replication/tests/test_compose_loss_integration.py::test_defaults_bit_exact_with_legacy_kwargs -xvs
pytest composer_replication/recipes/prime_rl/tests/test_composer_loss.py::test_returns_finite_scalar -xvs
```

---

### 10. `monarch` / `data-juicer` / `prime-rl` install (Wave 16)

**SYMPTOM.** `pip install -e ".[monarch]"`, `pip install -e ".[prime-rl]"`,
or `pip install -e ".[replaysim]"` fails immediately with a uv/pip
resolver error similar to:

```
× No solution found when resolving dependencies:
  ╰─▶ Because only monarch<=0.1.11 is available and
      composer-replication[monarch] depends on monarch>=0.4.1, we can
      conclude that composer-replication[monarch]'s requirements are
      unsatisfiable.
```

**DIAGNOSIS.** Three upstream packages the framework integrates with are
not currently pip-installable in their advertised versions:

1. **Meta's Monarch** is published on PyPI as
   `torchmonarch-nightly` (nightly wheels with platform constraints), not
   as `monarch`. The PyPI name `monarch` is unrelated to Meta's actor
   framework and tops out at `0.1.11`.
2. **Prime Intellect's prime-rl** is not registered on PyPI at all. It
   is published from source only.
3. **data-juicer** is not registered on PyPI under that exact name. The
   closest match (`py-data-juicer==1.0.0`) has broken transitive deps;
   newer `py-data-juicer` releases work but install ~150 transitive
   packages.

Wave 16 dropped all three extras from `pyproject.toml` rather than ship
unsatisfiable pins. The framework code paths that touch these libraries
import them lazily, so:
- `composer_replication.recipes.monarch` is a documentation skeleton
  that does NOT require monarch installed.
- `composer_replication.recipes.prime_rl.composer_loss` imports cleanly
  without prime-rl; the upstream parity test is `@skipif`-gated and the
  in-file shadow-parity test still verifies the loss formula
  independently.
- `composer_replication.replaysim.normalize.DJNormalizer(skip_dj=True)`
  works without `data_juicer`; only the full DJNormalizer code path
  needs it.

**FIX.** If you want any of these libraries' real functionality, install
from source alongside the framework:

```
# Meta Monarch (actor framework — see ADR-006)
pip install torchmonarch-nightly       # OR install from source:
# git clone https://github.com/meta-pytorch/monarch && cd monarch && pip install -e .

# Prime Intellect prime-rl (Recipe C — see ADR-006)
git clone https://github.com/PrimeIntellect-ai/prime-rl
cd prime-rl && pip install -e .

# data-juicer (replaysim normalization — see ADR-004)
git clone https://github.com/modelscope/data-juicer
cd data-juicer && pip install -e .
```

**VERIFICATION.** A fresh checkout install with all surviving extras
should succeed:

```
uv venv --clear
uv pip install -e ".[diloco,replay,replaysim,train,dev]"
source .venv/bin/activate
python -m pytest -q                    # baseline 176 passed / 8 skipped
```

If any of those extras fails to resolve, file a bug report — Wave 16
verified the full extras matrix installs from a clean venv on Python
3.11.

---

## How to file a bug report

If you've read the relevant section above and your problem persists,
file a bug. Include **all** sections of the template below — the most
common reason a maintainer can't repro is a missing piece of
environmental context.

```markdown
### What I expected vs what happened
(One paragraph.)

### Repro steps
1. ...
2. ...
3. ...

Minimal self-contained snippet (no `from my_local_thing import …`):

```python
# repro.py
from composer_replication import compose_loss
...
```

### Environment
- OS:                     (uname -a or `ver` on Windows)
- Python:                 (python --version)
- composer-replication:   (pip show composer-replication | head -3)
- torch:                  (python -c "import torch; print(torch.__version__)")
- torchft:                (python -c "import torchft; print(torchft.__version__)" || echo "n/a")
- transformers / trl:     (versions, or "not installed")
- data-juicer / fsspec:   (versions, or "not installed")
- s3fs / gcsfs / adlfs:   (versions if relevant)
- GPU:                    (nvidia-smi -L or "CPU only")
- Install method:         pip install -e . / wheel / other
- Extras installed:       [replay] [replaysim] [serverless] [dev]

### What you've already tried
- [ ] Read the relevant Failure Mode section of docs/TROUBLESHOOTING.md
      (which one: ___)
- [ ] Ran `pytest <relevant test path>` and confirmed those tests pass
- [ ] Ran the repro snippet in a fresh venv
- [ ] Confirmed it reproduces on Python 3.11 (if you were on 3.12 / 3.13)

### Logs
(Full traceback. If it's a wrong-loss-curve rather than an exception,
paste loss values for the first 10 steps and link any wandb/tb run.)

### Hypothesis
(Optional. If you have a guess at where the bug is, name the file +
line number. We'll look there first.)
```

A few rules:
- **Do not** paste API keys, AWS credentials, or HuggingFace tokens.
- **Do** include the failing test name if you've narrowed it to one.
- **Do** distinguish "never worked" from "regressed between commit X
  and Y." A regression-bisect goes straight to the front of the queue.
- **One bug per issue.** Multi-headed reports lose items in triage.

The Wave-14 surface area is large, but the test suite covers it
densely — every section above corresponds to a green test that proves
the fix worked.