Skip to content

Added flux demo - #3418

Merged
cehongwang merged 3 commits into
mainfrom
flux-demo
Jun 13, 2025
Merged

Added flux demo#3418
cehongwang merged 3 commits into
mainfrom
flux-demo

Conversation

@cehongwang

Copy link
Copy Markdown
Collaborator

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

  • 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)
  • This change requires a documentation update

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@github-actionsgithub-actionsBot added component: conversion Issues re: Conversion stage component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Feb 27, 2025
@cehongwang
cehongwang marked this pull request as draft February 27, 2025 00:39
Comment threadpy/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py Outdated
Comment threaddemo/flux_demo.py Outdated
@narendasan

Copy link
Copy Markdown
Collaborator

Can the app display the inference time, might be nice to have some stats rendered live as you generate

@cehongwang
cehongwangforce-pushed the flux-demo branch 4 times, most recently from 48a7c94 to 5a528f1CompareMarch 18, 2025 04:44
@github-actionsgithub-actionsBot added the component: tests Issues re: Tests label Mar 18, 2025
Comment threadpy/torch_tensorrt/runtime/_cudagraphs.py
@cehongwang
cehongwangforce-pushed the flux-demo branch 6 times, most recently from 361fb76 to 0aeea36CompareMarch 25, 2025 09:26
@cehongwang
cehongwang marked this pull request as ready for review March 26, 2025 07:53
@cehongwang
cehongwangforce-pushed the flux-demo branch 4 times, most recently from 9964674 to cfbc9eaCompareMarch 26, 2025 07:59
@cehongwang
cehongwangforce-pushed the flux-demo branch 4 times, most recently from 6d8023c to 0bcda12CompareMay 13, 2025 04:08
@cehongwang
cehongwangforce-pushed the flux-demo branch 3 times, most recently from 41139e9 to f536ac6CompareJune 2, 2025 20:38
Comment threadexamples/apps/flux-demo.py Outdated
Comment threadexamples/apps/NGRVNG.safetensors Outdated
Comment threadexamples/dynamo/torch_export_flux_dev.py Outdated
Comment threadexamples/dynamo/torch_export_flux_dev.py Outdated
Comment threadpy/torch_tensorrt/dynamo/_compiler.py
Comment threadpy/torch_tensorrt/dynamo/_refit.py
Comment threadpy/torch_tensorrt/dynamo/_refit.py Outdated
Comment threadpy/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py Outdated
Comment threadtools/perf/Flux/register_sdpa.py Outdated
from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
clean_up_graph_after_modifications,
)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

use them from examples

Comment threadtools/perf/Flux/flux_quantization.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I think we should avoid copying the whole model scripts for measuring perf. Try using the sys.path approach and importing the model and just a perf loop. something like

importsysimportossys.path.append(torchtrt_root+"examples/dynamo/apps")
fromflux_demoimport*model=<insertFLUXmodel (fp16orfp8) >results=measure_flux_perf(.... )

Comment threadexamples/apps/flux-demo.py Outdated
@cehongwang
cehongwangforce-pushed the flux-demo branch 6 times, most recently from 9acbee6 to 0d9a008CompareJune 12, 2025 23:00
@github-actionsgithub-actionsBot removed the component: lowering Issues re: The lowering / preprocessing passes label Jun 12, 2025

@github-actionsgithub-actionsBot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/quantize.py	2025-06-12 23:00:13.565607+00:00+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/quantize.py	2025-06-12 23:00:38.626299+00:00@@ -73,11 +73,10 @@
max_bound = 127
elif num_bits == 8 and exponent_bits == 4:
dtype = trt.DataType.FP8
max_bound = 448
-
axis = None
# int8 weight quantization is per-channel quantization(it can have one or multiple amax values)
if dtype == trt.DataType.INT8 and amax.numel() > 1:
# if the amax has more than one element, calculate the axis, otherwise axis value will be ignored
amax_init_shape = amax.shape
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py	2025-06-12 23:00:13.567607+00:00+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py	2025-06-12 23:00:39.100847+00:00@@ -98,16 +98,17 @@
class _TorchTensorRTConstantFolder(ConstantFolder): # type: ignore[misc]
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
def is_impure(self, node: torch.fx.node.Node) -> bool:
- # Set of known quantization ops to be excluded from constant folding. + # Set of known quantization ops to be excluded from constant folding.
# Currently, we exclude all quantization ops coming from modelopt library.
quantization_ops = {}
try:
- # modelopt import ensures torch.ops.tensorrt.quantize_op.default is registered + # modelopt import ensures torch.ops.tensorrt.quantize_op.default is registered
import modelopt.torch.quantization as mtq
+
assert torch.ops.tensorrt.quantize_op.default
quantization_ops.add(torch.ops.tensorrt.quantize_op.default)
except Exception as e:
pass
if quantization_ops and node.target in quantization_ops:

Comment threadexamples/apps/flux_demo.py Outdated
@github-actionsgithub-actionsBot added the component: lowering Issues re: The lowering / preprocessing passes label Jun 13, 2025

@github-actionsgithub-actionsBot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/quantize.py	2025-06-13 20:58:57.556250+00:00+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/quantize.py	2025-06-13 20:59:27.374709+00:00@@ -73,11 +73,10 @@
max_bound = 127
elif num_bits == 8 and exponent_bits == 4:
dtype = trt.DataType.FP8
max_bound = 448
-
axis = None
# int8 weight quantization is per-channel quantization(it can have one or multiple amax values)
if dtype == trt.DataType.INT8 and amax.numel() > 1:
# if the amax has more than one element, calculate the axis, otherwise axis value will be ignored
amax_init_shape = amax.shape

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signedcomponent: api [Python]Issues re: Python APIcomponent: conversionIssues re: Conversion stagecomponent: convertersIssues re: Specific op converterscomponent: dynamoIssues relating to the `torch.compile` or `torch._dynamo.export` pathscomponent: loweringIssues re: The lowering / preprocessing passescomponent: runtimecomponent: testsIssues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants

@cehongwang@narendasan@peri044@facebook-github-bot