') + ')', 'gi'); if (regex.test(text)) { found = true; var frag = document.createDocumentFragment(); var parts = text.split(regex); parts.forEach(function(part, i) { if (i % 2 === 0) { frag.appendChild(document.createTextNode(part)); } else { var span = document.createElement('span'); span.className = 'userscript-highlight'; span.textContent = part; frag.appendChild(span); } }); node.parentNode.replaceChild(frag, node); } }); } else if (node.nodeType === 1 && node.childNodes) { // element var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT']; if (!skipTags.includes(node.tagName)) { Array.from(node.childNodes).forEach(highlight); } } } highlight(document.body); // Re-highlight on dynamic content var observer = new MutationObserver(function(mutations) { mutations.forEach(function(m) { m.addedNodes.forEach(function(node) { if (node.nodeType === 1 || node.nodeType === 3) highlight(node); }); }); }); observer.observe(document.body, { childList: true, subtree: true }); })(); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + ', 'i'); if (__m === '*' || __re.test(location.href)) { // Strip utm_, fbclid, gclid, etc. from all links on page (function() { var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content', 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid', 'ref', 'ref_src', 'source', 'medium', 'campaign']; function cleanUrl(url) { try { var u = new URL(url, window.location.origin); var changed = false; trackingParams.forEach(function(p) { if (u.searchParams.has(p)) { u.searchParams.delete(p); changed = true; } }); return changed ? u.toString() : url; } catch (e) { return url; } } function cleanLinks() { document.querySelectorAll('a[href]').forEach(function(a) { var clean = cleanUrl(a.href); if (clean !== a.href) a.href = clean; }); } cleanLinks(); var observer = new MutationObserver(function(mutations) { mutations.forEach(function(m) { m.addedNodes.forEach(function(node) { if (node.nodeType === 1) { if (node.tagName === 'A') cleanLinks(); node.querySelectorAll('a[href]').forEach(function(a) { var clean = cleanUrl(a.href); if (clean !== a.href) a.href = clean; }); } }); }); }); observer.observe(document.body, { childList: true, subtree: true }); })(); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + ', 'i'); if (__m === '*' || __re.test(location.href)) { // Auto-enable theater mode on YouTube (function() { function tryTheater() { var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]'); if (btn && !btn.classList.contains('activated')) { btn.click(); } } // Try immediately tryTheater(); // Try after navigation (SPA) var lastUrl = location.href; setInterval(function() { if (location.href !== lastUrl) { lastUrl = location.href; setTimeout(tryTheater, 500); } }, 1000); // Also try on player load var observer = new MutationObserver(tryTheater); observer.observe(document.body, { childList: true, subtree: true }); })(); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + ', 'i'); if (__m === '*' || __re.test(location.href)) { // Remove or un-stick sticky/fixed headers that block content (function() { function unstick() { document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) { if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') { el.style.position = 'static'; el.style.top = 'auto'; el.style.zIndex = 'auto'; } }); } unstick(); var observer = new MutationObserver(unstick); observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] }); })(); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); })(); [PyTorch] Debug linear layer when saving original input and using debug quantizer by timmoon10 · Pull Request #1963 · NVIDIA/TransformerEngine · GitHub
Skip to content

[PyTorch] Debug linear layer when saving original input and using debug quantizer - #1963

Merged
timmoon10 merged 8 commits into
NVIDIA:mainfrom
timmoon10:debug-linear-original-input
Jul 22, 2025
Merged

[PyTorch] Debug linear layer when saving original input and using debug quantizer#1963
timmoon10 merged 8 commits into
NVIDIA:mainfrom
timmoon10:debug-linear-original-input

Conversation

@timmoon10

@timmoon10timmoon10 commented Jul 17, 2025•
edited
Loading

Copy link
Copy Markdown
Member

Description

#1865 introduced a failure in the distributed debug tests. The root cause is because we only all-gather the row-wise data for DebugQuantizedTensor:

TransformerEngine/transformer_engine/pytorch/distributed.py

Line 1394 in f8933bb

rowwise_total=gather_along_first_dim(rowwise, process_group, False, final_quantizer)[0]

However, if the linear layer is caching its original input tensor and requantizing in the backward pass, the correct behavior is to only quantize the column-wise data. This PR is a hacky workaround that only applies the debug quantizer to the gathered input tensor.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

  • Modify linear backward to avoid all-gathering debug tensor with only column-wise data

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@timmoon10
timmoon10 requested a review from pggPLJuly 17, 2025 23:43
@timmoon10
Workaround bugs with quantizing with only column-wise usage
6f66aa4
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
timmoon10force-pushed the debug-linear-original-input branch from 5387dd7 to 6f66aa4CompareJuly 17, 2025 23:52
@timmoon10

timmoon10 commented Jul 17, 2025

Copy link
Copy Markdown
MemberAuthor

/te-ci pytorch L1

@cyanguwacyanguwa mentioned this pull request Jul 18, 2025
Merged
13 tasks
@timmoon10
Remove unused imports
212a446
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
timmoon10 requested a review from ksivamanJuly 18, 2025 21:21
hxbai
hxbai reviewed Jul 20, 2025
View reviewed changes
Comment threadtransformer_engine/pytorch/module/linear.py
quantizer.set_usage(rowwise=True, columnwise=False)
else:
quantizer.set_usage(rowwise=False, columnwise=True)
quantizer.set_usage(rowwise=True, columnwise=True)

@hxbaihxbaiJul 20, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi Tim, why do we need the rowwise data here?

@timmoon10timmoon10Jul 21, 2025•
edited
Loading

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In principle we shouldn't need it, but I was running into issues where FP8 casts were failing without it. Actually, I don't think we need it once we skip the debug quantizer case.

@hxbaihxbaiJul 22, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this may be due to that the condition is modified. For (Float8Quantizer, Float8CurrentScalingQuantizer) cases, it doesn't support only quantizing the colwise data. But if backward_input_needs_gather is False, (Float8Quantizer, Float8CurrentScalingQuantizer) tensors would also go into the else path, which will cause the error. This hurts the performance of blockwise FP8 and MXFP8. Do you need me to create a fix PR for this, or will you fix this together with the DebugTensor?

@pggPL

pggPL commented Jul 21, 2025•
edited
Loading

Copy link
Copy Markdown
Collaborator

I think this line is source of error

TransformerEngine/transformer_engine/pytorch/distributed.py

Line 1408 in f8933bb

out_obj.rowwise_gemm_tensor=out_obj.rowwise_gemm_tensor

The rowwise/columnwise in DebugQuantizedTensors are the tensors used in gemms -> they can be both the same Float8Tensor object for example. And update_usage() does nothing in debug tensors currently.

@timmoon10

timmoon10 commented Jul 21, 2025

Copy link
Copy Markdown
MemberAuthor

@pggPL I tried changing

TransformerEngine/transformer_engine/pytorch/distributed.py

Line 1408 in f8933bb

out_obj.rowwise_gemm_tensor=out_obj.rowwise_gemm_tensor

to

out_obj.columnwise_gemm_tensor=out_obj.rowwise_gemm_tensor

However, the error reappeared when I applied the debug quantizer to the local input tensor.

For now, I think we should merge this as a quick bugfix and we can fix the edge cases for the debug tensor later.

ksivaman
ksivaman previously approved these changes Jul 21, 2025
View reviewed changes

@ksivamanksivaman left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@timmoon10

timmoon10 commented Jul 21, 2025

Copy link
Copy Markdown
MemberAuthor

/te-ci pytorch L1

@timmoon10
Workaround bugs with quantizing with only column-wise usage
d9c7dd1
FP8 does not support transpose-only cast.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
timmoon10 merged commit 315b47d into NVIDIA:mainJul 22, 2025
12 checks passed
@hxbaihxbai mentioned this pull request Jul 22, 2025
Merged
13 tasks
KshitijLakhani pushed a commit that referenced this pull request Jul 22, 2025
@timmoon10@pre-commit-ci@KshitijLakhani
[PyTorch] Debug linear layer when saving original input and using deb…
7ba6cd5
…ug quantizer (#1963)
* Debug linear layer when saving original input and using debug quantizer
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Workaround bugs with quantizing with only column-wise usage
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Remove unused imports
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Avoid unnecessary row-wise data
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Workaround bugs with quantizing with only column-wise usage
FP8 does not support transpose-only cast.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@OblynxOblynx mentioned this pull request Feb 8, 2026
Closed
@hungryGeek16hungryGeek16 mentioned this pull request May 31, 2026
Open
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ksivamanksivamanksivaman left review comments
@pggPLpggPLAwaiting requested review from pggPL
+1 more reviewer
@hxbaihxbaihxbai left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

2.6.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@timmoon10@pggPL@hxbai@ksivaman