|
2 | 2 | from typing import overload |
3 | 3 |
|
4 | 4 | import torch |
| 5 | +import torch.nn as nn |
5 | 6 | import torch.nn.functional as F |
| 7 | +from torch.distributed.fsdp import FSDPModule |
6 | 8 | from torch.nn import CrossEntropyLoss |
7 | 9 |
|
8 | 10 | from modalities.batch import InferenceResultBatch |
9 | 11 |
|
| 12 | +# PyTorch's default ignore index for cross-entropy loss |
| 13 | +IGNORE_INDEX = -100 |
| 14 | + |
10 | 15 |
|
11 | 16 | class Loss(ABC): |
12 | 17 | def __init__(self, tag: str): |
@@ -166,6 +171,190 @@ def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor: |
166 | 171 | return total_loss / flat_labels.numel() |
167 | 172 |
|
168 | 173 |
|
| 174 | +class ChunkedLMHeadCrossEntropyLoss(Loss): |
| 175 | + """Chunked lm_head + cross-entropy loss with per-chunk backward. |
| 176 | +
|
| 177 | + In contrast to ``ChunkedCLMCrossEntropyLoss`` (which chunks the loss math over |
| 178 | + already-materialized logits), this loss receives the *hidden states* from the model |
| 179 | + (see ``GPT2LLMConfig.return_hidden_states``) and applies the lm_head projection |
| 180 | + chunk-wise along the sequence dimension. The full ``[B, L, vocab_size]`` logits |
| 181 | + tensor is never materialized; peak memory for the head is ``O(B * L/num_chunks * V)``. |
| 182 | +
|
| 183 | + During training (``compute_and_backward``), each chunk's loss is backpropagated |
| 184 | + immediately through the lm_head so its autograd graph is freed before the next chunk |
| 185 | + runs. The gradients w.r.t. the hidden states are collected in a pre-allocated buffer |
| 186 | + and propagated through the transformer trunk with a single ``backward`` call at the |
| 187 | + end. This loss therefore OWNS the backward pass; the trainer must not call |
| 188 | + ``loss.backward()`` again. |
| 189 | +
|
| 190 | + Under FSDP2, the lm_head must be its own ``fully_shard`` unit (see |
| 191 | + ``ModelFactory.get_fsdp2_wrapped_model(..., wrap_lm_head_separately=True)``). The |
| 192 | + lm_head weight is kept unsharded across all chunks (avoiding repeated all-gathers) |
| 193 | + and its gradient reduce-scatter is coalesced into the last chunk's backward. |
| 194 | + Weight tying is not supported in that setting, since wte and lm_head would then |
| 195 | + belong to different FSDP units. |
| 196 | +
|
| 197 | + Args: |
| 198 | + target_key: Key to access labels in the ``InferenceResultBatch``. |
| 199 | + prediction_key: Key to access the hidden states in the ``InferenceResultBatch``. |
| 200 | + num_chunks: Number of chunks to split the sequence dimension into. The sequence |
| 201 | + length must be divisible by this. Defaults to 8 (as in torchtitan). |
| 202 | + use_compile: Apply ``torch.compile`` to the per-chunk cross-entropy function. |
| 203 | + Defaults to ``True``. |
| 204 | + tag: Loss tag used for logging. |
| 205 | + """ |
| 206 | + |
| 207 | + def __init__( |
| 208 | + self, |
| 209 | + target_key: str, |
| 210 | + prediction_key: str, |
| 211 | + num_chunks: int = 8, |
| 212 | + use_compile: bool = True, |
| 213 | + tag: str = "ChunkedLMHeadCrossEntropyLoss", |
| 214 | + ): |
| 215 | + super().__init__(tag) |
| 216 | + if num_chunks < 1: |
| 217 | + raise ValueError(f"num_chunks must be >= 1, got {num_chunks}") |
| 218 | + self.target_key = target_key |
| 219 | + self.prediction_key = prediction_key |
| 220 | + self.num_chunks = num_chunks |
| 221 | + self.lm_head: nn.Module | None = None |
| 222 | + self._loss_fn = torch.compile(_ce_loss_sum) if use_compile else _ce_loss_sum |
| 223 | + |
| 224 | + def bind_lm_head(self, model: nn.Module) -> None: |
| 225 | + """Binds the (possibly FSDP-wrapped) lm_head module of the given model to this loss. |
| 226 | +
|
| 227 | + Must be called after the model is fully wrapped (FSDP etc.) and before the first |
| 228 | + loss computation. The trainer and evaluator call this lazily. |
| 229 | + """ |
| 230 | + self.lm_head = model.get_submodule("transformer.lm_head") |
| 231 | + if isinstance(self.lm_head, FSDPModule): |
| 232 | + # FSDP2's default backward-prefetch heuristic picks its target from a global, |
| 233 | + # cross-module post-forward-call order. Since lm_head.forward() is invoked |
| 234 | + # num_chunks times per step, interleaved with its own per-chunk backward calls |
| 235 | + # (forward(chunk0) -> backward(chunk0) -> forward(chunk1) -> ...), that ordering |
| 236 | + # no longer reflects real data dependencies: the heuristic ends up re-gathering |
| 237 | + # unrelated modules (e.g., the last transformer block, which no chunk backward |
| 238 | + # actually needs, since each hidden_chunk is a detached leaf) or lm_head's own |
| 239 | + # earlier forward call. Overriding the prefetch target to lm_head itself makes |
| 240 | + # this a no-op (unshard() returns immediately when already unsharded, which it |
| 241 | + # is throughout the chunk loop -- see compute_and_backward), eliminating the |
| 242 | + # wasted re-gathers. An empty list would NOT disable prefetching here: FSDP |
| 243 | + # treats an empty override list the same as no override (falls back to the |
| 244 | + # default heuristic), so a self-reference is required. |
| 245 | + self.lm_head.set_modules_to_backward_prefetch([self.lm_head]) |
| 246 | + |
| 247 | + def _chunk(self, t: torch.Tensor) -> list[torch.Tensor]: |
| 248 | + seq_len = t.shape[1] |
| 249 | + if seq_len % self.num_chunks != 0: |
| 250 | + raise ValueError( |
| 251 | + f"Sequence length {seq_len} must be divisible by num_chunks {self.num_chunks}." |
| 252 | + ) |
| 253 | + return list(t.tensor_split(self.num_chunks, dim=1)) |
| 254 | + |
| 255 | + def _chunk_ce_sum(self, hidden_chunk: torch.Tensor, label_chunk: torch.Tensor) -> torch.Tensor: |
| 256 | + logits = self.lm_head(hidden_chunk) |
| 257 | + return self._loss_fn(logits.reshape(-1, logits.size(-1)), label_chunk.reshape(-1).long()) |
| 258 | + |
| 259 | + def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor: |
| 260 | + """Forward-only chunked loss (used for evaluation, typically under ``torch.no_grad()``).""" |
| 261 | + if self.lm_head is None: |
| 262 | + raise RuntimeError( |
| 263 | + "lm_head is not set. Call bind_lm_head(model) before computing " |
| 264 | + "ChunkedLMHeadCrossEntropyLoss." |
| 265 | + ) |
| 266 | + hidden_states = forward_batch.get_predictions(self.prediction_key) |
| 267 | + labels = forward_batch.get_targets(self.target_key).to(hidden_states.device) |
| 268 | + |
| 269 | + num_valid = (labels != IGNORE_INDEX).sum().clamp(min=1) |
| 270 | + total_loss = hidden_states.new_zeros(()) |
| 271 | + for hidden_chunk, label_chunk in zip(self._chunk(hidden_states), self._chunk(labels)): |
| 272 | + total_loss = total_loss + self._chunk_ce_sum(hidden_chunk, label_chunk) |
| 273 | + return total_loss / num_valid |
| 274 | + |
| 275 | + def compute_and_backward( |
| 276 | + self, hidden_states: torch.Tensor, labels: torch.Tensor, grad_scale: float = 1.0 |
| 277 | + ) -> torch.Tensor: |
| 278 | + """Computes the loss chunk-wise and runs the full backward pass. |
| 279 | +
|
| 280 | + Args: |
| 281 | + hidden_states: Model output of shape [B, L, D] (after lm_head_norm), requires grad. |
| 282 | + labels: Target token ids of shape [B, L]; ``IGNORE_INDEX`` entries are masked out. |
| 283 | + grad_scale: Factor multiplied into the backpropagated loss only (e.g., |
| 284 | + 1/gradient_accumulation_steps). The returned loss value is unscaled. |
| 285 | +
|
| 286 | + Returns: |
| 287 | + The detached mean loss over non-ignored tokens. |
| 288 | + """ |
| 289 | + if self.lm_head is None: |
| 290 | + raise RuntimeError( |
| 291 | + "lm_head is not set. Call bind_lm_head(model) before computing " |
| 292 | + "ChunkedLMHeadCrossEntropyLoss." |
| 293 | + ) |
| 294 | + if not hidden_states.requires_grad: |
| 295 | + raise RuntimeError( |
| 296 | + "compute_and_backward requires hidden_states with requires_grad=True. " |
| 297 | + "Use __call__ for evaluation." |
| 298 | + ) |
| 299 | + |
| 300 | + labels = labels.to(hidden_states.device) |
| 301 | + hidden_chunks = [c.detach().requires_grad_() for c in self._chunk(hidden_states)] |
| 302 | + label_chunks = self._chunk(labels) |
| 303 | + |
| 304 | + num_valid = (labels != IGNORE_INDEX).sum().clamp(min=1) |
| 305 | + grad_buffer = torch.zeros_like(hidden_states) |
| 306 | + total_loss = hidden_states.new_zeros(()) |
| 307 | + |
| 308 | + # Keep the lm_head weight unsharded across all chunks (one all-gather instead of |
| 309 | + # num_chunks) and coalesce the gradient reduce-scatter into the last chunk's backward. |
| 310 | + is_fsdp2 = isinstance(self.lm_head, FSDPModule) |
| 311 | + if is_fsdp2: |
| 312 | + self.lm_head.set_reshard_after_forward(False) |
| 313 | + self.lm_head.set_reshard_after_backward(False) |
| 314 | + # The chunk backwards are intermediate backwards of the same training step. By |
| 315 | + # default FSDP2 treats every backward() as the last one: its root post-backward |
| 316 | + # callback then finalizes ALL param groups, clearing the global post_forward_order |
| 317 | + # and every transformer block's _post_forward_indices -- the bookkeeping that |
| 318 | + # backward prefetching relies on. The subsequent trunk backward would then run |
| 319 | + # with prefetching disabled, exposing every block's all-gather (~10 ms GPU idle |
| 320 | + # per block, +~280 ms/step measured for the 8B model at micro-batch size 2). |
| 321 | + # Marking the chunk backwards as non-last skips that finalization. The flag lives |
| 322 | + # on the state context shared by the entire FSDP tree, so setting it through the |
| 323 | + # lm_head handle also governs the root callback. |
| 324 | + self.lm_head.set_is_last_backward(False) |
| 325 | + |
| 326 | + offset = 0 |
| 327 | + for chunk_idx, (hidden_chunk, label_chunk) in enumerate(zip(hidden_chunks, label_chunks)): |
| 328 | + is_last = chunk_idx == self.num_chunks - 1 |
| 329 | + if is_fsdp2: |
| 330 | + self.lm_head.set_requires_gradient_sync(is_last) |
| 331 | + if is_last: |
| 332 | + self.lm_head.set_reshard_after_backward(True) |
| 333 | + |
| 334 | + chunk_loss = self._chunk_ce_sum(hidden_chunk, label_chunk) |
| 335 | + # Backward through the lm_head only (hidden_chunk is a detached leaf); frees |
| 336 | + # this chunk's activations before the next chunk's forward. |
| 337 | + (chunk_loss * (grad_scale / num_valid)).backward() |
| 338 | + |
| 339 | + chunk_len = hidden_chunk.shape[1] |
| 340 | + grad_buffer[:, offset : offset + chunk_len] = hidden_chunk.grad |
| 341 | + hidden_chunk.grad = None |
| 342 | + offset += chunk_len |
| 343 | + total_loss = total_loss + chunk_loss.detach() |
| 344 | + |
| 345 | + if is_fsdp2: |
| 346 | + self.lm_head.set_reshard_after_forward(True) |
| 347 | + self.lm_head.reshard() |
| 348 | + # The trunk backward is the true last backward of this step: let the root |
| 349 | + # callback finalize all param groups and clear the prefetch bookkeeping. |
| 350 | + self.lm_head.set_is_last_backward(True) |
| 351 | + |
| 352 | + # Single backward through the transformer trunk with the accumulated gradient. |
| 353 | + hidden_states.backward(gradient=grad_buffer) |
| 354 | + |
| 355 | + return total_loss / num_valid |
| 356 | + |
| 357 | + |
169 | 358 | def nce_loss( |
170 | 359 | embedding1: torch.Tensor, embedding2: torch.Tensor, device: torch.device, is_asymmetric: bool, temperature: float |
171 | 360 | ) -> torch.Tensor: |
|
0 commit comments