Aller au contenu
login
arrow_backRetour aux issues
MakazhanAlpamys/Soup #425

stream_layers: true crashes on pre-Ampere (T4) GPUs — fp16 GradScaler receives bf16 LoRA adapters

ecoDébutant bug help wanted

descriptionDescription

## Summary Running the `proof-4gb.ipynb` §5 headline (or any `stream_layers: true` config) on a free Colab T4 (sm_75, Turing) crashes on the **first** training step: ``` NotImplementedError: "_amp_foreach_non_finite_check_and_unscale_cuda" not implemented for 'BFloat16' ``` ## Environment - Hardware: free-tier Colab, Tesla T4 (sm_75), 15.36 GB VRAM, 12 GB RAM, 2 vCPU - Python 3.12.13, torch 2.11.0+cu128, CUDA 12.8 - Soup 0.73.2 from `main` (`pip install "soup-cli[train] @ git+https://github.com/MakazhanAlpamys/Soup.git"`), `torchao` uninstalled per the notebook - `NousResearch/Meta-Llama-3.1-8B-Instruct`, `quantization: 4bit`, `stream_layers: true`, `stream_buffers: 2`, LoRA r=8 ## Repro Follow `notebooks/proof-4gb.ipynb` through §5. Crash at the first `clip_grad_norm_`: ``` File ".../soup_cli/trainer/sft.py", line 1521, in train self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) File ".../transformers/trainer.py", line 2715, in _inner_training_loop _grad_norm = self.accelerator.clip_grad_norm_(...) File ".../accelerate/accelerator.py", line 3006, in clip_grad_norm_ self.unscale_gradients() File ".../accelerate/accelerator.py", line 2944, in unscale_gradients self.scaler.unscale_(opt) File ".../torch/amp/grad_scaler.py", line 358, in unscale_ optimizer_state["found_inf_per_device"] = self._unscale_grads_(...) File ".../torch/amp/grad_scaler.py", line 294, in _unscale_grads_ torch._amp_foreach_non_finite_check_and_unscale_(...) NotImplementedError: "_amp_foreach_non_finite_check_and_unscale_cuda" not implemented for 'BFloat16' ``` ## Root cause (diagnosed) `pick_mixed_precision` correctly returns `fp16` on a T4 (`cc 7.5 < BF16_MIN_CC`), so `TrainingArguments(fp16=True, bf16=False)` and the GradScaler are fp16. But peft's `get_peft_model` creates the LoRA adapters in the **base checkpoint's dtype** — bf16 for Llama-3.1 — independent of the stream dtype. On Ampere that is consistent (bf16 stream + bf16 mixed precision); on a pre-Ampere card the stream and mixed precision are fp16 while all **128 trainable adapter params stay bf16**. The fp16 `GradScaler.unscale_()` then dispatches the `foreach` kernel against bf16 gradients and raises. Probes confirmed the two things that stop an in-place cast from saving you: 1. `materialize_meta_adapters` (which re-inits to fp32) never corrects this — the adapters are **not** `meta` (peft allocates real tensors), so the `if not param.is_meta` guard skips them. 2. A cast added inside `build_streamed_model` *does* fire (`adapters_off_dtype 128 -> 0`), but `accelerator.prepare()` re-casts the adapters back to bf16 before the trainer runs. ## Proposed fix Cast trainable params to fp32 at the latest point — after `accelerator.prepare()`, before the lazily-created optimizer. In `trainer/sft.py`, immediately before `self.trainer.train(...)`: ```python if getattr(self.trainer.args, "fp16", False) and not getattr(self.trainer.args, "bf16", False): import torch for p in self.trainer.model.parameters(): if p.requires_grad and p.dtype == torch.bfloat16: p.data = p.data.to(torch.float32) ``` fp32 matches the author's own `materialize_meta_adapters` choice and provides fp32 master weights for stable LoRA updates. ## Verification With the fix, the identical config completes: 7 steps, peak **2.86 GB** (vs the recorded 2.91 GB), adapter written with **128/128 tensors non-zero**. ## Related measurement (fills a gap flagged in `benchmarks/run-t4-colab-free-tier.md`) That record notes the §4 streamed-vs-resident comparison "produced no captured output ... recorded as unrun." I ran it on the T4: `torch.equal(streamed, resident)` is **False** (max |diff| ≈ 3.3 on logits of magnitude up to ~38). This is fp16 non-associativity, not a correctness defect: argmax agrees **98.2%**, and against an fp32 reference both models sit within fp16 tolerance (resident err 0.40, streamed err 0.37). The streamed and resident
codeOuvre sur GitHub