rename selfgen_only -> selfgen, add asr_gender (388180)

README.md

@@ -10,24 +10,41 @@ Whisper-large-v3 encoder (frozen) → concat+MLP adapter → Qwen3-4B-Instruct-2
 Only the adapter and the LoRA weights are trained, so each checkpoint is ~147 MB; the base
 models are downloaded from their own repos at load time.
 
-| folder | training data | test-clean WER | test-clean gender |
+Code, configs and the full recipe: <https://github.com/kehanlu/interspeech-tutorial>
+
+| folder | training data | test-clean ASR | test-clean gender |
 |---|---|---|---|
-| `selfgen_only` | 281k self-generated conversational responses (no task labels) | 3.93 (best prompt + post-processing) | 98.24 |
-| `selfgen_mix` | the same + 20% ASR + 5% gender SFT | 2.45 | 98.82 |
+| `asr_gender` | 281k ASR + a fresh 30% of the gender rows each epoch | **1.81** WER | **98.85** |
+| `selfgen` | 281k self-generated conversational replies, no task labels | 3.93 WER | 98.24 |
 
-The self-generation targets were written by Qwen3-4B given only the transcript and the
-speaker's gender: `<audio>{transcription} (Gender: {gender})</audio>` + "The audio is a
-passage read aloud from a book. Respond directly as a natural conversation partner. Do not
-mention the audio, the transcription, or the speaker attributes."
+`asr_gender` is ordinary task SFT, and it matches whisper-large-v3 on ASR (1.89) while also
+answering the gender question. It is the baseline.
 
-`selfgen_only` has never seen a transcription instruction: it can transcribe, but needs an
-explicit prompt ("Transcribe the speech word for word. Output only the transcription …") and
-sometimes answers in the training-time metadata format. `selfgen_mix` follows the plain
-prompts directly.
+`selfgen` is the interesting one: its targets were written by Qwen3-4B given only the
+transcript and the speaker's gender, so **it has never seen a transcription or a gender
+label as a training target**. The self-generation prompt was
 
-## Usage
+    <audio>{transcription} (Gender: {gender})</audio>
+
+    The audio is a passage read aloud from a book. Respond directly as a natural
+    conversation partner. Do not mention the audio, the transcription, or the speaker
+    attributes.
+
+It can still do both tasks, but only if the prompt leaves room for a short answer. Asked the
+way `asr_gender` was trained ("Transcribe the speech into text") it replies with an essay
+about the passage and scores 52.35 WER; asked for a format it reaches 3.93:
 
-See the Colab notebook in the tutorial repo, or:
+| prompt | ASR |
+|---|---|
+| `Transcribe the speech into text` | 52.35 → 16.54 after clean-up |
+| `Transcribe the speech word for word. Output only the transcription, with no explanation, in this format:\nAnswer: "<transcription>"` | 8.94 → **3.93** |
+
+Gender goes 81.87 → **98.24** the same way, with "The audio is a passage read aloud from a
+book. Is the speaker male or female? Answer with one word." The clean-up is the rule-based
+`postprocess()` in the tutorial repo's `example/evaluate/evaluate_asr.py`; it is a no-op on
+`asr_gender`, which already answers with a bare transcript.
+
+## Usage
 
 ```python
 from huggingface_hub import hf_hub_download, snapshot_download
@@ -37,9 +54,12 @@ code = snapshot_download("kehanlu/interspeech-tutorial", allow_patterns="code/*"
 sys.path.insert(0, f"{code}/code")
 from inference import SpeechLLMForInference
 
-ckpt = hf_hub_download("kehanlu/interspeech-tutorial", "selfgen_mix/model.ckpt")
+ckpt = hf_hub_download("kehanlu/interspeech-tutorial", "selfgen/model.ckpt")
 pipe = SpeechLLMForInference.from_checkpoint(ckpt, dtype=torch.float16)   # float16 for a Colab T4
 print(pipe.generate([{"role": "user",
                       "content": "<audio><|AUDIO|></audio>\n\nTranscribe the speech into text",
                       "audios": [{"audio": "sample.flac"}]}]))
 ```
+
+A few LibriSpeech dev-clean clips are in `samples/` with their transcripts in
+`samples/samples.json`.

asr_gender/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b29485636775546d881df994666f4f6a2e68bc7548b75303a13279948e9bb19
+size 146924946

code/data.py

@@ -1,7 +1,6 @@
 """
 Data pipeline for the minimal SpeechLLM.
 
-The one idea to hold on to while reading this file: **the LLM never sees audio.**
 It sees a plain text token sequence in which every ``<|AUDIO|>`` has been replaced
 by exactly as many placeholder tokens as that audio will occupy after encoding.
 The model later overwrites those placeholders' *embeddings* with real speech
@@ -24,7 +23,7 @@ Manifest format (one JSON object per line)::
       "target": "A wooden door being knocked on."
     }
 
-`audio_filepath` is relative to `data_root`. `target` may also be called
+`audio_filepath` is relative to `audio_root`. `target` may also be called
 `response` (both appear in the DeSTA3 manifests). Extra keys are ignored.
 """
 
@@ -45,32 +44,25 @@ logger = logging.getLogger(__name__)
 SAMPLE_RATE = 16000
 
 
-# --------------------------------------------------------------------------
 # How many tokens will one audio occupy?
 #
 # Whisper's feature extractor always pads to 30 s, so we cannot use the padded
 # length -- a 5 s clip would reserve 25 s worth of slots. We ask the feature
 # extractor for an attention mask over the mel frames instead, and push the true
-# length through the same conv arithmetic the model uses.
-#
-# These formulas mirror Qwen2AudioEncoder._get_feat_extract_output_lengths.
-# --------------------------------------------------------------------------
+# length through the same downsampling arithmetic the model uses.
 
 def whisper_frames(mel_len: int) -> int:
     """Whisper's encoder conv2 has stride 2. 3000 mel frames -> 1500 positions."""
     return (mel_len - 1) // 2 + 1
 
 
-def conv1d_out(length: int, kernel: int = 3, stride: int = 2, padding: int = 1) -> int:
-    return (length + 2 * padding - kernel) // stride + 1
-
+def audio_token_count(mel_len: int, n_downsample: int) -> int:
+    """mel frames -> number of LLM token slots this audio needs.
 
-def audio_token_count(mel_len: int, n_downsample_layers: int) -> int:
-    """mel frames -> number of LLM token slots this audio needs."""
-    length = whisper_frames(mel_len)
-    for _ in range(n_downsample_layers):
-        length = conv1d_out(length)
-    return length
+    The connector concatenates every `n_downsample` adjacent frames and right-pads to
+    a multiple of that, so the count is ceil(frames / n_downsample), rounding up.
+    """
+    return -(-whisper_frames(mel_len) // n_downsample)
 
 
 def expand_audio_locator(tokens: List[str], audio_locator: str,
@@ -140,25 +132,15 @@ def manifest_name(path: str) -> str:
 
 class AudioTextDataset(Dataset):
     """Lazy jsonl reader with per-manifest mixing ratios.
-
-    The DeSTA3 manifests run to hundreds of MB each, so we build a byte-offset
-    index up front and read one line at a time instead of holding the parsed
-    rows in memory.
-
-    `ratios[i]` scales manifest i: 0.3 keeps a random 30% of its rows, 2.5 keeps two
-    full copies plus a random 50%. The random part is drawn from (seed, epoch,
-    manifest) only -- never the rank -- so every DDP rank builds the same index
-    and DistributedSampler shards it consistently. Call `resample(epoch)` to draw a
-    fresh subset for a new epoch; the length does not change.
     """
 
-    def __init__(self, manifest_filepaths, data_root: str, ratios=None, seed: int = 0):
-        self.data_root = data_root
+    def __init__(self, manifest_filepaths, audio_root: str, ratios=None, seed: int = 0):
+        self.audio_root = audio_root
         self.paths = [resolve_manifest(p) for p in manifest_filepaths]
         self.ratios = list(ratios) if ratios is not None else [1.0] * len(self.paths)
         assert len(self.ratios) == len(self.paths)
         self.seed = seed
-        self.rows = []           # per manifest: [(file_idx, byte_offset), ...]
+        self.rows = []  # per manifest: [(file_idx, byte_offset), ...]
         for file_idx, path in enumerate(self.paths):
             rows = []
             with open(path, "rb") as f:
@@ -168,7 +150,7 @@ class AudioTextDataset(Dataset):
                         rows.append((file_idx, offset))
                     offset += len(line)
             self.rows.append(rows)
-        self._handles = {}       # opened lazily, per dataloader worker
+        self._handles = {}  # opened lazily, per dataloader worker
         self.resample(epoch=0, log=True)
 
     @property
@@ -177,7 +159,7 @@ class AudioTextDataset(Dataset):
         return any(r != math.floor(r) for r in self.ratios)
 
     def resample(self, epoch: int, log: bool = False):
-        self.index = []          # (file_idx, byte_offset)
+        self.index = []  # (file_idx, byte_offset)
         for file_idx, (rows, ratio) in enumerate(zip(self.rows, self.ratios)):
             whole = math.floor(ratio)
             n_extra = round((ratio - whole) * len(rows))
@@ -206,7 +188,7 @@ class AudioTextDataset(Dataset):
         assert target, f"row {i} has neither `response` nor `target`"
 
         audios = [{"audio_filepath": resolve_audio_filepath(
-                       os.path.join(self.data_root, audio["audio_filepath"]))}
+                       os.path.join(self.audio_root, audio["audio_filepath"]))}
                   for audio in row["audios"]]
         return {"messages": row["messages"], "target": target, "audios": audios}
 
@@ -226,7 +208,7 @@ class Collator:
     placeholder_token: str = "<|vision_pad|>"
     max_seq_length: int = 1024
     max_audio_seconds: float = 30.0
-    n_downsample_layers: int = 1
+    n_downsample: int = 2
     # inference: build the prompt only, with no answer appended and no labels.
     # Same slot arithmetic as training -- that is the whole point of reusing this
     # class rather than rebuilding the expansion in the demo script.
@@ -253,7 +235,7 @@ class Collator:
             return_tensors="pt", return_attention_mask=True,
         )
         mel_lengths = features["attention_mask"].sum(-1).tolist()
-        audio_lengths = [audio_token_count(n, self.n_downsample_layers) for n in mel_lengths]
+        audio_lengths = [audio_token_count(n, self.n_downsample) for n in mel_lengths]
 
         # 2. text -> token ids, with <|AUDIO|> expanded to placeholders
         context_ids, target_ids = [], []
@@ -280,7 +262,7 @@ class Collator:
             # placeholder can never be re-tokenized into something else
             ctx = self.tokenizer.convert_tokens_to_ids(tokens)
             if self.for_generation:
-                tgt = []                       # nothing to condition on; the model writes it
+                tgt = []  # nothing to condition on; the model writes it
             else:
                 tgt = self.tokenizer.encode(row["target"], add_special_tokens=False)
                 tgt = tgt + [self.tokenizer.eos_token_id]
@@ -288,7 +270,7 @@ class Collator:
             assert len(ctx) < self.max_seq_length, (
                 f"prompt alone is {len(ctx)} tokens (max_seq_length="
                 f"{self.max_seq_length}); shorten the audio or raise the limit")
-            tgt = tgt[:self.max_seq_length - len(ctx)]      # only ever truncate the answer
+            tgt = tgt[:self.max_seq_length - len(ctx)]  # only ever truncate the answer
 
             context_ids.append(ctx)
             target_ids.append(tgt)
@@ -315,6 +297,6 @@ class Collator:
             "attention_mask": attention_mask,
             "labels": labels,
             "input_features": features["input_features"],
-            "audio_lengths": audio_lengths,          # per audio, after downsampling
-            "start_positions": shifted_starts,       # per audio, (row, first placeholder)
+            "audio_lengths": audio_lengths,  # per audio, after downsampling
+            "start_positions": shifted_starts,  # per audio, (row, first placeholder)
         }

code/inference.py

@@ -8,20 +8,8 @@ Inference for the minimal SpeechLLM -- the tutorial's demo entry point.
     # score a manifest: prints reference vs hypothesis side by side
     python inference.py --ckpt exp/run/checkpoints/last.ckpt \
         --manifest data/Librispeech-dev-test/test-clean_asr.jsonl \
-        --data-root /home/u8915687/work/data/audios2 --limit 20
-
-The architecture is read back out of the checkpoint's `hyper_parameters`, never
-hand-written here. That matters more than it looks: `SpeechLLMModule` loads with
-`strict=False`, and the connector's downsampling stack is an `nn.Sequential`
-whose keys are positional (`downsample.0.*`, `downsample.2.*`). Build the model
-with the wrong `n_downsample_layers` and those tensors become "missing keys",
-which `strict=False` silently drops -- you get a randomly-initialised connector,
-no error, and fluent nonsense out. `_load_trainable` below turns that into a
-hard failure instead.
-
-Unlike DeSTA3 there is no VAD and no ASR decoder in the loop: DeSTA3 runs Whisper's
-decoder to produce a transcript, and needs VAD to decide which clips are speech
-at all. Here the LLM only sees the speech features, exactly as in training.
+        --audio-root /path/to/data --limit 20
+
 """
 
 import argparse
@@ -57,7 +45,7 @@ class SpeechLLMForInference:
         self.model, self.tokenizer = model, tokenizer
         self.collator, self.device, self.dtype = collator, device, dtype
 
-    # -- loading ----------------------------------------------------------
+    # -- loading --
     @classmethod
     def from_checkpoint(cls, ckpt_path, device=None, dtype=None):
         device = device or ("cuda" if torch.cuda.is_available() else "cpu")
@@ -66,9 +54,12 @@ class SpeechLLMForInference:
         dtype = dtype or getattr(torch, hp["dtype"])
         logger.info("checkpoint %s (epoch %s, step %s)",
                     os.path.basename(ckpt_path), ckpt.get("epoch"), ckpt.get("global_step"))
-        logger.info("  llm=%s encoder=%s adapter=%s n_downsample=%d frozen_encoder=%s",
-                    hp["llm_id"], hp["encoder_id"], hp.get("adapter_type", "cnn"),
-                    hp["n_downsample_layers"], hp.get("freeze_encoder", False))
+        # n_downsample used to be n_downsample_layers, an exponent: 1 meant 2 frames
+        # per token. Map the old key so earlier checkpoints still load.
+        n_downsample = hp.get("n_downsample") or 2 ** hp["n_downsample_layers"]
+        logger.info("  llm=%s encoder=%s n_downsample=%d frozen_encoder=%s",
+                    hp["llm_id"], hp["encoder_id"],
+                    n_downsample, hp.get("freeze_encoder", False))
 
         # rebuilt exactly as SpeechLLMModule.__init__ does, or the ids shift
         tokenizer = AutoTokenizer.from_pretrained(hp["llm_id"])
@@ -79,9 +70,8 @@ class SpeechLLMForInference:
 
         model = SpeechLLM(
             llm_id=hp["llm_id"], encoder_id=hp["encoder_id"],
-            n_downsample_layers=hp["n_downsample_layers"], dtype=dtype,
-            # .get(): checkpoints from before these options were cnn + trained encoder
-            adapter_type=hp.get("adapter_type", "cnn"),
+            n_downsample=n_downsample, dtype=dtype,
+            # .get(): older checkpoints predate these options
             adapter_hidden_dim=hp.get("adapter_hidden_dim", 0),
             freeze_encoder=hp.get("freeze_encoder", False), use_lora=True,
             lora_rank=hp["lora_rank"], lora_alpha=hp["lora_alpha"],
@@ -95,7 +85,7 @@ class SpeechLLMForInference:
             tokenizer=tokenizer, feature_extractor=feature_extractor,
             audio_locator=hp["audio_locator"], placeholder_token=hp["placeholder_token"],
             max_seq_length=hp["max_seq_length"], max_audio_seconds=hp["max_audio_seconds"],
-            n_downsample_layers=hp["n_downsample_layers"],
+            n_downsample=n_downsample,
             for_generation=True,
         )
         return cls(model, tokenizer, collator, device, dtype)
@@ -117,7 +107,7 @@ class SpeechLLMForInference:
         logger.info("  restored %d trainable tensors (%.1fM params)", len(state),
                     sum(v.numel() for v in state.values()) / 1e6)
 
-    # -- generation -------------------------------------------------------
+    # -- generation --
     def _to_rows(self, conversations):
         """DeSTA3-style messages -> the row dicts `Collator` consumes."""
         rows = []
@@ -140,7 +130,7 @@ class SpeechLLMForInference:
     @torch.no_grad()
     def generate(self, conversations, max_new_tokens=200, do_sample=False, **generation_kwargs):
         if conversations and isinstance(conversations[0], dict):
-            conversations = [conversations]          # a single conversation
+            conversations = [conversations]  # a single conversation
         batch = self.collator(self._to_rows(conversations))
 
         batch["input_ids"] = batch["input_ids"].to(self.device)
@@ -163,7 +153,7 @@ class SpeechLLMForInference:
                 self.tokenizer.batch_decode(generated, skip_special_tokens=True)]
 
 
-# ------------------------------------------------------------------ main ---
+# -- main --
 
 def main():
     ap = argparse.ArgumentParser(description=__doc__,
@@ -171,7 +161,7 @@ def main():
     ap.add_argument("--ckpt", required=True)
     ap.add_argument("--audio", nargs="+", help="one or more audio files")
     ap.add_argument("--manifest", help="jsonl to run over instead of --audio")
-    ap.add_argument("--data-root", default="")
+    ap.add_argument("--audio-root", default="")
     ap.add_argument("--limit", type=int, default=10, help="rows to take from --manifest")
     ap.add_argument("--prompt", default=DEFAULT_PROMPT)
     ap.add_argument("--batch-size", type=int, default=4)
@@ -183,7 +173,7 @@ def main():
     pipe = SpeechLLMForInference.from_checkpoint(args.ckpt, device=args.device)
 
     if args.audio:
-        items = [(os.path.join(args.data_root, p), None) for p in args.audio]
+        items = [(os.path.join(args.audio_root, p), None) for p in args.audio]
     else:
         items = []
         with open(args.manifest) as f:
@@ -191,7 +181,7 @@ def main():
                 if len(items) >= args.limit:
                     break
                 row = json.loads(line)
-                items.append((os.path.join(args.data_root,
+                items.append((os.path.join(args.audio_root,
                                            row["audios"][0]["audio_filepath"]),
                               row.get("response") or row.get("target")))
 

code/modeling.py

@@ -4,15 +4,6 @@ A minimal SpeechLLM: Whisper encoder + connector + LLM with LoRA.
     audio ──► Whisper encoder (trained or frozen) ──► connector (trained) ──┐
                                                                             ├──► LLM + LoRA ──► text
     text  ────────────────────────────────────────► embedding table ────────┘
-
-Two connectors (`adapter_type`), both 50 fps -> 25 fps per downsampling step:
-    cnn         strided Conv1d + GELU, then LayerNorm + Linear
-    concat_mlp  concatenate adjacent frames, then LayerNorm + Linear + GELU + Linear
-
-A frozen LLM does *not* make training cheap: the speech features enter at layer 0,
-so gradients still flow back through every LLM layer to reach the connector, and
-every LLM activation has to be kept for the backward pass. Turn on
-`gradient_checkpointing` if that hurts.
 """
 
 import logging
@@ -24,74 +15,36 @@ from transformers import AutoModelForCausalLM, WhisperModel
 logger = logging.getLogger(__name__)
 
 
-class SpeechConnector(nn.Module):
-    """Whisper hidden states → LLM-width embeddings.
-
-    Takes the encoder's last hidden state and applies two steps:
-
-    1. **Strided Conv1d downsampling.** Whisper emits 50 frames/s, which is far
-       too many tokens to hand an LLM. One stride-2 convolution brings that to
-       25 frames/s -- 40 ms per token, the same rate as Qwen2-Audio. Each extra
-       layer halves the rate again.
-
-    2. **Linear projection to the LLM's width** (from Qwen2-Audio, whose entire
-       connector is ``AvgPool1d(2)`` + ``Linear``).
-    """
-
-    def __init__(self, d_encoder: int, d_llm: int, n_downsample_layers: int = 1):
-        super().__init__()
-        layers = []
-        for _ in range(n_downsample_layers):
-            layers += [nn.Conv1d(d_encoder, d_encoder, kernel_size=3, stride=2, padding=1),
-                       nn.GELU()]
-        self.downsample = nn.Sequential(*layers)
-
-        self.proj = nn.Sequential(nn.LayerNorm(d_encoder), nn.Linear(d_encoder, d_llm))
-
-    def forward(self, hidden_states):                          # (B, T, d_encoder)
-        hidden_states = self.downsample(hidden_states.transpose(1, 2)).transpose(1, 2)
-        return self.proj(hidden_states)                        # (B, T', d_llm)
-
-
 class ConcatMLPConnector(nn.Module):
     """Whisper hidden states → LLM-width embeddings by frame concatenation.
-
-    Every `2 ** n_downsample_layers` adjacent frames are concatenated into one
-    feature (d_encoder * k), which a two-layer MLP maps to the LLM's width. This is
-    the SLAM-ASR / LLaVA-style projector: no convolution, all mixing happens in the
-    MLP. The sequence is right-padded to a multiple of k, so a length-L input gives
-    ceil(L / k) tokens -- the same count as the CNN, so `data.audio_token_count`
-    serves both.
     """
 
-    def __init__(self, d_encoder: int, d_llm: int, n_downsample_layers: int = 1,
+    def __init__(self, d_encoder: int, d_llm: int, n_downsample: int = 2,
                  hidden_dim: int = 0):
         super().__init__()
-        self.k = 2 ** n_downsample_layers
+        self.n_downsample = n_downsample
         hidden_dim = hidden_dim or d_llm
         self.proj = nn.Sequential(
-            nn.LayerNorm(d_encoder * self.k),
-            nn.Linear(d_encoder * self.k, hidden_dim),
+            nn.LayerNorm(d_encoder * n_downsample),
+            nn.Linear(d_encoder * n_downsample, hidden_dim),
             nn.GELU(),
             nn.Linear(hidden_dim, d_llm),
         )
 
     def forward(self, hidden_states):                          # (B, T, d_encoder)
         B, T, D = hidden_states.shape
-        pad = (-T) % self.k
+        k = self.n_downsample
+        pad = (-T) % k
         if pad:
             hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad))
-        stacked = hidden_states.reshape(B, (T + pad) // self.k, D * self.k)
+        stacked = hidden_states.reshape(B, (T + pad) // k, D * k)
         return self.proj(stacked)                              # (B, ceil(T/k), d_llm)
 
 
-ADAPTERS = {"cnn": SpeechConnector, "concat_mlp": ConcatMLPConnector}
-
-
 class SpeechLLM(nn.Module):
     def __init__(self, llm_id: str, encoder_id: str,
-                 n_downsample_layers: int = 1, dtype=torch.bfloat16,
-                 adapter_type: str = "cnn", adapter_hidden_dim: int = 0,
+                 n_downsample: int = 2, dtype=torch.bfloat16,
+                 adapter_hidden_dim: int = 0,
                  gradient_checkpointing: bool = False, cache_dir=None,
                  freeze_encoder: bool = True, use_lora: bool = False,
                  lora_rank: int = 32, lora_alpha: int = 64, lora_dropout: float = 0.05,
@@ -104,25 +57,15 @@ class SpeechLLM(nn.Module):
         self.encoder = WhisperModel.from_pretrained(
             encoder_id, torch_dtype=dtype, cache_dir=cache_dir).encoder
 
-        assert adapter_type in ADAPTERS, f"adapter_type must be one of {list(ADAPTERS)}"
-        extra = {"hidden_dim": adapter_hidden_dim} if adapter_type == "concat_mlp" else {}
-        self.connector = ADAPTERS[adapter_type](
+        self.connector = ConcatMLPConnector(
             d_encoder=self.encoder.config.d_model,
             d_llm=self.llm.config.hidden_size,
-            n_downsample_layers=n_downsample_layers,
-            **extra,
+            n_downsample=n_downsample,
+            hidden_dim=adapter_hidden_dim,
         )
-        logger.info("adapter %s: %.1fM params", adapter_type,
+        logger.info("connector: %.1fM params",
                     sum(p.numel() for p in self.connector.parameters()) / 1e6)
 
-        # --- what trains ---------------------------------------------------
-        # The connector always does. The LLM trains only through LoRA adapters,
-        # if at all. The encoder is a switch: leaving it frozen is much cheaper
-        # than it looks, because with every encoder parameter frozen and
-        # `input_features` not requiring grad, PyTorch builds no graph for the
-        # encoder at all -- its activations are never stored. Unfreezing it
-        # turns 32 Whisper layers back into stored activations, which costs far
-        # more memory than the 635M parameters themselves.
         for param in self.llm.parameters():
             param.requires_grad = False
 
@@ -134,8 +77,7 @@ class SpeechLLM(nn.Module):
                 target_modules=list(lora_target_modules), bias="none",
                 task_type="CAUSAL_LM",
             )
-            # unwrap back to a plain module: we call the LLM directly (with
-            # inputs_embeds) rather than through the PeftModel wrapper
+
             self.llm = get_peft_model(self.llm, lora_config).base_model.model
             n_lora = sum(p.numel() for n, p in self.llm.named_parameters() if "lora_" in n)
             assert n_lora > 0, f"LoRA matched no modules in {lora_target_modules}"
@@ -146,25 +88,12 @@ class SpeechLLM(nn.Module):
         for param in self.encoder.parameters():
             param.requires_grad = not freeze_encoder
 
-        # Standard mixed precision: frozen weights stay in `dtype`, everything that
-        # trains is stored in float32 and autocast does the bf16 compute. bf16 has
-        # only 7 fraction bits (~0.8% between neighbouring values), so AdamW's
-        # small per-step updates would round to zero on bf16 weights.
-        # The connector and LoRA are created in float32 already; the encoder is
-        # loaded in `dtype`, so cast it when it trains.
         if not freeze_encoder:
             self.encoder.float()
 
         if gradient_checkpointing:
-            # use_reentrant=False matters under DDP: the reentrant implementation
-            # runs the recomputation outside autograd's view, so DDP can mark a
-            # gradient ready twice and abort with "Expected to mark a variable
-            # ready only once". The non-reentrant path composes with DDP.
             ckpt_kwargs = {"use_reentrant": False}
             self.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs=ckpt_kwargs)
-            # deliberately no enable_input_require_grads(): the spliced speech
-            # features already make inputs_embeds require grad, and that hook
-            # would make it a leaf, breaking the splice (see build_inputs_embeds)
             if not freeze_encoder:
                 self.encoder.gradient_checkpointing_enable(
                     gradient_checkpointing_kwargs=ckpt_kwargs)
@@ -188,10 +117,6 @@ class SpeechLLM(nn.Module):
 
     def encode_audio(self, input_features):
         """mel features → speech embeddings in the LLM's width.
-
-        We run the encoder over the full 30 s window including padding, and crop
-        afterwards. Whisper is *trained* that way, so this is its native regime;
-        masking the padding would be off-distribution for a frozen encoder.
         """
         input_features = input_features.to(self.encoder.conv1.weight.dtype)
         return self.connector(self.encoder(input_features).last_hidden_state)
@@ -199,12 +124,6 @@ class SpeechLLM(nn.Module):
     def build_inputs_embeds(self, input_ids, speech_features, audio_lengths, start_positions):
         """Overwrite each audio's placeholder embeddings with its speech features."""
         embed = self.llm.get_input_embeddings()
-        # The embedding table is frozen, so `embed(input_ids)` comes back as a
-        # leaf with requires_grad=False -- and writing in-place into a leaf that
-        # requires grad is illegal, which is exactly what happens once gradient
-        # checkpointing turns that flag on. Clone to get a non-leaf we may write
-        # into. Splicing the speech features in is what makes the result require
-        # grad, which is also what checkpointing needs from the LLM's input.
         inputs_embeds = embed(input_ids).clone()
 
         for k, (row, start) in enumerate(start_positions):

selfgen/model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbea96fded8610bd5d66b5576c3d8c5635b2ba3e063896bd32c9ec562910c548
+size 146924882