Repository navigation
Conversation
transformers loads the tensors a model lists in `_keep_in_fp32_modules` (under fp16) and `_keep_in_fp32_modules_strict` (under fp16 and bf16) in fp32, because they overflow or lose precision in half precision: T5's `wo` projection (huggingface/transformers#20683, added for FlanT5-XL), the MoE routing bias `e_score_correction_bias` of DeepSeek-V3/GLM-4.5 and others. `InferenceEngine._convert_to_dtype` then calls `self.module.half()` / `.bfloat16()` on the whole model, which casts those tensors along with everything else. For FlanT5 in fp16 that is the divergence from Hugging Face reported in deepspeedai#3177: `wo` overflows and generation goes wrong with no error. Record the fp32 tensors the model's lists name (reusing `keep_in_fp32_pattern` from the ZeRO-3 fix in deepspeedai#8699), cast the module, and restore them, so `init_inference` leaves them as transformers loaded them. A listed tensor that arrives in a lower precision is upcast, as transformers' loader and deepspeedai#8699 do. Measured on CPU with google/flan-t5-base in fp16: before this change `init_inference(model, dtype=torch.float16)` generates different text than the same model in Hugging Face; after it the outputs are identical. Fixes deepspeedai#3177 Signed-off-by: Udaya Tejas <udayatejas2004@gmail.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 003707e00f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| keep = self._tensors_to_keep_in_fp32(config.dtype) | ||
| self.module.half() | ||
| self._restore_fp32_tensors(keep) |
There was a problem hiding this comment.
When a T5 model has already been loaded in fp16 on an accelerator (with its many wo matrices left in fp32), keep retains every fp32 storage while self.module.half() allocates and installs a half copy of every one; none of those temporary copies is released until _restore_fp32_tensors runs after the whole cast. Peak memory is therefore the final mixed-precision model plus an extra fp16 copy of all protected matrices, which can OOM large T5 variants even when the final model fits. Cast nonmatching tensors selectively, and validate peak memory on a representative model.
AGENTS.md reference: AGENTS.md:L37-L37
Useful? React with 👍 / 👎.
| for name, param in engine.module.named_parameters(): | ||
| expected = torch.float32 if name.endswith("wo.weight") else torch.float16 | ||
| assert param.dtype == expected, name |
There was a problem hiding this comment.
When fp16 is available, this test stops at parameter dtypes and never invokes the wrapped T5. The reported regression is incorrect generation, so an activation/weight dtype mismatch or output divergence caused by the mixed fp32 wo path would pass every assertion here. Add a minimal forward or generation and compare its logits or tokens with an independently derived or Hugging Face reference; inference-level contract changes require device-executed integration coverage.
AGENTS.md reference: AGENTS.md:L81-L82
Useful? React with 👍 / 👎.
This branch has not been deployed
1 participant
Fixes #3177
Problem
transformers loads the tensors a model lists in
_keep_in_fp32_modules(under fp16) and_keep_in_fp32_modules_strict(under fp16 and bf16) in fp32, because they overflow or lose precision inhalf precision. T5/MT5/UMT5/LongT5/SwitchTransformers list their
woprojection(huggingface/transformers#20683, added in Dec 2022 for exactly the FlanT5-XL fp16 case), and 17 MoE
model files (DeepSeek-V3, GLM-4.5, ...) list the routing bias
e_score_correction_biasunder both dtypes.InferenceEngine._convert_to_dtype(deepspeed/inference/engine.py:521-524) then callsself.module.half()/self.module.bfloat16()on the whole model and casts those tensors with everythingelse. With
mp_size=1and no kernel injection (the reporter's call) nothing else ininit_inferencetouches the model, so this cast is the entire difference from Hugging Face.
Impact: silent-wrong-result. Who reaches it / triggered by: anyone calling
deepspeed.init_inference(model, dtype=torch.float16)(fp16 is also the config's defaultdtype) ordtype=torch.bfloat16on a transformers model with one of those lists, e.g. the script in #3177. What isobserved: generations that differ from the same model run in transformers, with no warning or error.
Measured here on CPU with
google/flan-t5-baseloaded withdtype=torch.float16(transformers 5.17.0,torch 2.14.0+cpu; the CPU accelerator's fp16 gate was bypassed in the script so the cast path runs), on
the two sentences from the issue, greedy,
max_length=32:The
.half()line dates from the 2021 inference release (ed3de0c), before transformers introduced thelists, so this is drift rather than a decision. The training side gained
keep_in_fp32_patternfor thesame lists in #8699 (ZeRO-3); the inference engine was left out.
Fix
Before casting, collect the fp32 tensors whose full names match the model's lists (reusing
keep_in_fp32_pattern(module, "auto", dtype)from #8699, which applies the strict list under bf16 and bothlists under fp16, the way transformers does), cast the module, and put those tensors back. A listed tensor
that arrives in a lower precision is upcast, as transformers' loader and #8699 do. The kept tensors are
the ones transformers already runs with under fp16, and the HF forward passes already cast activations to
the weight dtype for them, so the module ends in the state transformers itself produces. Kernel injection
and AutoTP run after this point and are unchanged.
Testing
New
tests/unit/inference/test_inference_keep_in_fp32.py: a toy module with both lists, parametrized overfp16/bf16 (fp16 skips where the accelerator lacks it), plus a real
T5ForConditionalGenerationbuilt froma tiny config under fp16.
Negative control (revert only
deepspeed/inference/engine.pytoupstream/master, keep the new test), ontest_init_inference_keeps_listed_tensors_in_fp32[dtype1](bf16, the one this CPU can run):Restored:
1 passed. The fp16 cases pass with the fix and fail without it when the fp16 gate is bypassedlocally; on CPU CI they skip and will run on the GPU jobs.
test_inference_config.py: 9 passed, 1pre-existing failure (
test_cuda_graph_with_kernel_inject_raisesfails identically on unmodifiedupstream/masterbecause this CPU has no fp16).Execution hardware: CPU only (no GPU available to me), torch 2.14.0+cpu.
pre-commit run --filescleanon both changed files.
Signed-off-by: Udaya Tejas udayatejas2004@gmail.com