[New model] PersonaPlex (full-duplex speech model, based on Moshi)
Author: un1876
Pull request: no
Labels: name: New model

### Model description

# [Moshi] Modernize for PersonaPlex: modular + processor + WeightConverter + #36928 fix

# What does this PR do?

**Our research team:** @b-re-w, @un1876, @jookite, @capyjin, @haesol-shin

We are using [NVIDIA PersonaPlex](https://huggingface.co/nvidia/personaplex-7b-v1) in an ongoing research project on full-duplex spoken dialogue. In the course of that research — at [latentforge/transformers-moshi](https://github.com/latentforge/transformers-moshi) — Moshi has been modernized to match current transformers conventions, and PersonaPlex support is being built on top of it. Rather than keeping that work on a fork, I would like to contribute it here. The overall work:

1. Converts Moshi to the modular format (`modular_moshi.py`)
2. Adds `MoshiProcessor` and feature extraction
3. Migrates checkpoint conversion to `WeightConverter`, removing `convert_moshi_transformers.py`
4. Fixes the double-shift in loss computation reported in #36928 — a bug we ran into ourselves and resolved the same way, with that report as our reference
5. **Builds toward PersonaPlex support — the reason for 1–4.** This is the model our research runs on: we have been studying the released checkpoint and building the implementation to fit the actual PersonaPlex model. That work is in progress on the fork and will land in a follow-up, but it is what shaped 1–4, so I would rather put it in front of you here than surface it after the fact.

**The code in this PR is items 1–4; item 5 is what they are for.** (1) and (2) apply conventions that `kyutai_speech_to_text` — the model that inherits from Moshi — already follows, so this also lets the downstream model sit on a parent that is no longer the odd one out. (3) follows the v5 dynamic weight loading API. (4) is the bug reported in #36928 — we hit it ourselves in our own training runs, and resolved it the same way with that report as our reference; measurements below.

Related: while looking through open issues, I found that Moshi appears in #39588's list of models still lacking a modular file (`moshi (2512 lines in modeling file)`) — if this contribution lands, it looks like that entry would be resolved along the way.

For (5), see [Next step: PersonaPlex](#next-step-personaplex) below — what I have already found in the checkpoint, the direction I am building toward, and the questions I would like your input on now rather than once the code is written.

## Why: Moshi is the only model in its own family still on the legacy pattern

Moshi is the base of this family — `kyutai_speech_to_text` inherits directly from it:

```python
# src/transformers/models/kyutai_speech_to_text/modular_kyutai_speech_to_text.py
from ..moshi.modeling_moshi import MoshiModel, MoshiPreTrainedModel

class KyutaiSpeechToTextPreTrainedModel(MoshiPreTrainedModel): ...
class KyutaiSpeechToTextModel(MoshiModel): ...
```

Yet the child was added with the current conventions while the base it derives from was never migrated:

| | `kyutai_speech_to_text` | `moshi` (before this PR) |
|---|---|---|
| `modular_*.py` | yes | no |
| `processing_*.py` | yes | no |
| `feature_extraction_*.py` | yes | no |

This shows up directly in user-facing code. Kyutai STT exposes a single processor:

```python
processor = KyutaiSpeechToTextProcessor.from_pretrained(model_id)
inputs = processor(ds[0]["audio"]["array"])
```

while Moshi requires assembling the feature extractor and the tokenizer by hand — this is the current docs example:

```python
feature_extractor = AutoFeatureExtractor.from_pretrained("kyutai/moshiko-pytorch-bf16")
tokenizer = AutoTokenizer.from_pretrained("kyutai/moshiko-pytorch-bf16")
user_input_values = feature_extractor(raw_audio=audio_sample, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt")
```

Moshi is a text + audio model and the processor convention has a first-class `audio` argument, so a `MoshiProcessor` seemed like the natural fit. `KyutaiSpeechToTextProcessor` is the closest reference in the family, so I followed its shape rather than inventing a new one.

## The four changes in this PR

### 1. Modular conversion

Pure refactor, no behavior change. Verified with `check_modular_conversion.py` and the existing test suites (below).

### 2. MoshiProcessor and feature extraction

Wraps the feature extractor and the tokenizer into a single processor, following `KyutaiSpeechToTextProcessor`. This is the only part of the PR that adds new public API.

### 3. Checkpoint conversion via WeightConverter

Registers Moshi in `conversion_mapping.py` so the original checkpoint loads directly through `from_pretrained`. Since the conversion then happens at load time, `convert_moshi_transformers.py` is no longer needed and is removed in this PR.

Two things I would like your call on here:

- This one touches `conversion_mapping.py`, which is a central file rather than something contained in the model folder. If you would rather keep the standalone script and skip this part, I am happy to drop it — it is a separate commit.
- I scoped this to Moshi only. `kyutai_speech_to_text` still has its own `convert_*.py` script, so the two are inconsistent on this axis after this PR.

### 4. Double-shift in loss computation (fixes #36928)

This is the bug reported in #36928 — we ran into it ourselves in our own training runs, and resolved it the same way, with that report as our reference. `MoshiForCausalLM.forward` states its own intent in a comment, then does something else:

```python
# Shift so that tokens < n predict n      <- the stated goal
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
...
loss = self.loss_function(shift_logits, shift_labels, ...)   # positional
```

Because the class name contains `ForCausalLM`, `LOSS_MAPPING` resolves `self.loss_function` to `ForCausalLMLoss`, which carries the *same* comment and shifts again when its `shift_labels` argument is `None`. The two positional arguments land on `logits`/`labels`, so the `shift_labels` switch never engages and the shift happens twice.

**The experiment asks one question:** does the loss the model reports score `logits[t]` against `labels[t+1]` — what the comment promises — or against `labels[t+2]`?

**Data and model.** A tiny randomly initialized Moshi and one fixed input. Random data is enough because the alignment is a property of the code, not of the data:

```python
torch.manual_seed(42)   # fixes the weights
model = MoshiForCausalLM(MoshiConfig(vocab_size=99, hidden_size=32, num_hidden_layers=2,
                                     num_attention_heads=2, num_key_value_heads=2, ffn_dim=64)).eval()
torch.manual_seed(7)    # fixes the input
input_ids = torch.randint(0, 99, (1, 8));  labels = input_ids.clone()
```

**Procedure.** The same script runs twice; the only thing that changes in between is `modeling_moshi.py` (`git checkout`). Each run records the loss the model reports, plus two candidate answers computed from the same logits with plain `F.cross_entropy` — no transformers loss code, so the code under test never grades itself:

```python
with torch.no_grad():
    out = model(input_ids=input_ids, labels=labels)             # what the model reports
ref_1 = F.cross_entropy(logits[0, :-1], labels[0, 1:]).item()   # scored against labels[t+1]
ref_2 = F.cross_entropy(logits[0, :-2], labels[0, 2:]).item()   # scored against labels[t+2]
```

**Result** — raw output of the two runs:

```
before: {"loss": 4.550253, "match_ref1": false, "match_ref2": true,  "logits_sum": 2.3901}
after:  {"loss": 4.551546, "match_ref1": true,  "match_ref2": false, "logits_sum": 2.3901}
```

The identical `logits_sum` confirms both runs saw the same weights and input, so the flip can only come from the file: the reported loss reproduces one candidate to six decimals — before the fix it *is* the 2-shift value, after it *is* the 1-shift value. (The two candidates sit close together only because an untrained model scores ≈ `ln(99) ≈ 4.6` either way; the evidence is the exact match, not the gap. `LlamaForCausalLM` and `GPT2LMHeadModel` under the same harness both match `ref_1`.)

**Severity, on data whose answer is known by construction.** Build logits so that position `t` is certain about `labels[t+1]`, with a 60-logit gap between the right token and every other one. Any correctly aligned loss must read ≈ 0 here — no reference formula needed:

```python
logits = torch.full((1, 6, 20), -30.0)
for t in range(5):
    logits[0, t, labels[0, t + 1]] = 30.0
```

| logits built so that… | before | after |
|---|---|---|
| position `t` predicts `labels[t+1]` | **60.0** | **0.0** |
| position `t` predicts `labels[t+2]` | **0.0** | **60.0** |

Chance level here is `ln(20) ≈ 3.0`; 60.0 is the price of being maximally confident and wrong. The pre-fix code scores 0.0 for predicting two tokens ahead and 60.0 for predicting the next one — the opposite of what its own comment asks for.

<details>
<summary>Self-contained reproduction script (~25 lines)</summary>

```python
# pip install torch, then run once per file state:
#   python repro_shift.py                                        # with the fix
#   git checkout src/transformers/models/moshi/modeling_moshi.py
#   python repro_shift.py                                        # without it
import json

import torch
import torch.nn.functional as F

from transformers import MoshiConfig
from transformers.models.moshi.modeling_moshi import MoshiForCausalLM

torch.manual_seed(42)   # fixes the weights
model = MoshiForCausalLM(MoshiConfig(vocab_size=99, hidden_size=32, num_hidden_layers=2,
                                     num_attention_heads=2, num_key_value_heads=2, ffn_dim=64)).eval()
torch.manual_seed(7)    # fixes the input
input_ids = torch.randint(0, 99, (1, 8))
labels = input_ids.clone()

with torch.no_grad():
    out = model(input_ids=input_ids, labels=labels)
logits = out.logits.float()
loss = out.loss.item()
ref_1 = F.cross_entropy(logits[0, :-1], labels[0, 1:]).item()   # scored against labels[t+1]
ref_2 = F.cross_entropy(logits[0, :-2], labels[0, 2:]).item()   # scored against labels[t+2]
print(json.dumps({"loss": round(loss, 6),
                  "match_ref1": abs(loss - ref_1) < 1e-5,
                  "match_ref2": abs(loss - ref_2) < 1e-5,
                  "logits_sum": round(logits.sum().item(), 4)}))
```

</details>

**The fix** removes the manual shift (which duplicates what `ForCausalLMLoss` already does, including the `.float()` upcast, the flatten and the `.to(device)`) and passes the tensors by keyword, matching `modeling_llama.py`:

```python
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
```

This is a separate commit and it does change behavior, so it is easy to drop if you would rather land it via #36928 itself. Note the keyword form also makes the original mistake structurally impossible. If you would like the alignment locked in, I am happy to add the `ref_1` identity as a regression test in this PR.

## These changes also affect kyutai_speech_to_text

Because `kyutai_speech_to_text` imports from `moshi.modeling_moshi`, converting Moshi to modular is not a Moshi-only change. I ran the Kyutai suite as part of this PR, and there are no regressions:

| Check | Result |
|---|---|
| `ruff` (moshi + kyutai) | clean |
| modular consistency (moshi + kyutai) | exit 0 |
| `check_copies` | exit 0 |
| Moshi tests | 200 passed |
| Kyutai tests | 105 passed |

`kyutai_speech_to_text` is the only model in the library that inherits from Moshi, so it is the full blast radius of changes 1–3. Change 4 is scoped to `MoshiForCausalLM.forward`, which Kyutai does not inherit (it subclasses `MoshiModel` / `MoshiPreTrainedModel`), so its loss path is untouched.

## Would you like this split into separate PRs?

The commits are separated by concern:

- modular conversion — pure refactor, no behavior change
- processor / feature extraction — new public API
- weight conversion — `WeightConverter` + removal of the convert script
- loss fix — behavior change, fixes #36928

I kept them together because they are one "bring Moshi up to current conventions" pass, but I am happy to split them — in particular the last two are self-contained and can be dropped or landed separately.

## Next step: PersonaPlex

This is item (5) from the top — the model our research is built on, and the reason changes 1–4 exist. Support is already in development on the fork; the code will come in a follow-up, but I would like the approach settled here, in this thread. Fitting it to the released checkpoint is how I ran into the `dep_q` difference below.

**About the model.** Per its model card, PersonaPlex finetunes Moshi (Moshiko weights) and adds a "Hybrid System Prompt" mechanism that combines text-based role conditioning with audio-based voice prompting.

- Paper: [arXiv:2602.06053](https://arxiv.org/abs/2602.06053)
- Weights: [nvidia/personaplex-7b-v1](https://huggingface.co/nvidia/personaplex-7b-v1)
- Original code: [github.com/NVIDIA/personaplex](https://github.com/NVIDIA/personaplex)

**What I have verified so far.** The released PersonaPlex checkpoint trains its depth transformer with `dep_q=16` — it predicts both the agent and the user audio streams — whereas the released Moshi checkpoint uses `dep_q=8` (agent stream only). So there is at least a config-level difference beyond the prompt-handling logic.

**The direction I am currently building toward.** The `kyutai_speech_to_text` pattern: its own model folder with a `modular_personaplex.py` inheriting from Moshi, plus a `WeightConverter` so the original NVIDIA checkpoint loads directly.

## Questions

1. **Splitting the PRs** — if we proceed, should the Moshi modernization and the PersonaPlex support go as separate PRs?

2. **The #36928 loss fix** — is it OK to fold that fix into this PR?
   - That PR currently has merge conflicts and has been inactive for a while.
   - If the fix lands here, could #36928 then be closed as resolved — with credit to its author?

3. **PersonaPlex shape** — a lighter path also seems possible: since PersonaPlex is ultimately a finetune of Moshi, it might work without a new model class — loading the checkpoint into the existing Moshi classes, with `dep_q=16` as a config value and the system-prompt handling in the processor. Which of the two would fit the library better?

This modernization came out of our research rather than being an end in itself. It would be great to see it land upstream instead of staying on a fork — and if it does, we hope it makes a useful contribution to the library. I am happy to take it wherever you want it to go: split it, drop any part of it, or rework things to match a convention I have missed.

## AI tool usage

Per the contributing guide: AI coding assistants were used while preparing this PR, for both the implementation and the write-up. Every change was reviewed and tested by the authors — the reproduction script and the loss measurements above were run and verified by us, and the test results (Moshi 200, Kyutai 105) were produced locally.

## Who can review?

@eustlb (Moshi PRs / contributed Kyutai Speech-To-Text), @ylacombe (contributed Moshi)


### Open source status

- [x] The model implementation is available
- [x] The model weights are available

### Provide useful links for the implementation

- Implementation: https://github.com/NVIDIA/personaplex (MIT)
- Weights: https://huggingface.co/nvidia/personaplex-7b-v1
- Paper: https://arxiv.org/abs/2602.06053

Comments: