Fix deepgemm on multiple devices#47323
Conversation
| # DeepGEMM linear off). Also temporarily skipped under ``torch.compile`` — DeepGEMM's | ||
| # per-token cast calls ``pack_ue8m0_to_int`` which has data-dependent bit-twiddling that | ||
| # dynamo can't guard. TODO: remove the ``is_torchdynamo_compiling`` gate once the upstream | ||
| # ``pack_ue8m0_to_int`` is rewritten to be FakeTensor-friendly; the Triton fallback is | ||
| # dynamo-friendly today via its ``@triton_op`` registration. |
There was a problem hiding this comment.
this was already fixed
There was a problem hiding this comment.
Ah nice, is it already also the correct version we use and we just forgot?
There was a problem hiding this comment.
yes i already submitted the fix as part of the dsv4 static cache PR but that never got merged (perf not so great for now)
| block_size: list[int] | None = None, | ||
| bias: torch.Tensor | None = None, | ||
| activation_scale: torch.Tensor | None = None, | ||
| output_dtype: torch.dtype | None = None, |
There was a problem hiding this comment.
this was just bad because it defaulted to None and deepgemm and fp8 don't behave the same on None output dtype anyways
There was a problem hiding this comment.
Should we add * at least to be not BC in case anyone used the interface? Similar to others, it is slightly breaking because an output dtype is no longer possible
Maybe we could default to input dtype on None instead?
There was a problem hiding this comment.
defaulted to input.dtype but i think we should deprecate it, in general linears don't really interface an output dtype, see torch's linear https://docs.pytorch.org/docs/2.13/generated/torch.nn.functional.linear.html
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| block_size: list[int] | None = None, | ||
| bias: torch.Tensor | None = None, | ||
| activation_scale: torch.Tensor | None = None, | ||
| output_dtype: torch.dtype | None = None, |
There was a problem hiding this comment.
Should we add * at least to be not BC in case anyone used the interface? Similar to others, it is slightly breaking because an output dtype is no longer possible
Maybe we could default to input dtype on None instead?
| # DeepGEMM linear off). Also temporarily skipped under ``torch.compile`` — DeepGEMM's | ||
| # per-token cast calls ``pack_ue8m0_to_int`` which has data-dependent bit-twiddling that | ||
| # dynamo can't guard. TODO: remove the ``is_torchdynamo_compiling`` gate once the upstream | ||
| # ``pack_ue8m0_to_int`` is rewritten to be FakeTensor-friendly; the Triton fallback is | ||
| # dynamo-friendly today via its ``@triton_op`` registration. |
There was a problem hiding this comment.
Ah nice, is it already also the correct version we use and we just forgot?
| class FP8Linear(nn.Linear): | ||
| # Set True at load when the model spans >1 CUDA device in one process; DeepGEMM's | ||
| # context-bound kernels corrupt across devices (see `quantizer_finegrained_fp8.py`). | ||
| _deepgemm_disabled = False |
There was a problem hiding this comment.
nit, would set on init either way, no?
There was a problem hiding this comment.
no only when we have than one device do we update it, i can remove this but will have to do getattr and it will be there sometimes and sometimes not so i thought it could make sense to add it like the _can_compile_fullgraph and other capability flags
| class FP8Experts(nn.Module): | ||
| # Set True at load when the model spans >1 CUDA device in one process; DeepGEMM's | ||
| # context-bound kernels corrupt across devices (see `quantizer_finegrained_fp8.py`). | ||
| _deepgemm_disabled = False |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: finegrained_fp8 |
CI recapDashboard: View test results in Grafana |
3 participants
What does this PR do?
Fixes # (issue)
Code Agent Policy
The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. These often are low-quality, or fix extremely minor issues that occur rarely or never in practice.
As a result, we're instituting a rule that first-time contributors should not use code agents to submit PRs or issues.
We'd also ask autonomous "OpenClaw"-like agents not to open any PRs or issues.
Issues/PRs from first-time contributors that violate this rule will probably just be closed without review, and we
might block you, especially if you open more than one or appear to be deliberately ignoring this. We especially do not
want new contributors to jump in on random issues to contribute an agent-written fix. This creates lots of noise
for reviewers and other users and will almost certainly get you blocked.
For more information, please read
CONTRIBUTING.md.Before submitting
Pull Request checks?
to it if that's the case.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.