Skip to content

[tests] fix torchao tests - #14258

Open
sayakpaul wants to merge 2 commits into
mainfrom
torchao-fixes
Open

[tests] fix torchao tests#14258
sayakpaul wants to merge 2 commits into
mainfrom
torchao-fixes

Conversation

@sayakpaul

Copy link
Copy Markdown
Member

No description provided.

@sayakpaul
sayakpaul requested review from DN6 and dg845 July 21, 2026 07:54
@github-actions github-actions Bot added tests size/S PR with diff < 50 LOC labels Jul 21, 2026
Comment on lines +872 to +880
# version=2 int4 defaults to the "plain" packing format, which routes through the
# fbgemm/mslk Int4Tensor kernels. Pin the packing format to the tinygemm
# (_convert_weight_to_int4pack) path on CUDA and plain_int32 on Intel XPU so the tests
# don't require those extra kernels to be installed.
if config_name == "Int4WeightOnlyConfig":
if torch_device == "xpu":
config_kwargs["int4_packing_format"] = "plain_int32"
elif torch_device == "cuda":
config_kwargs["int4_packing_format"] = "tile_packed_to_4d"

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.

Is this change intended mainly for XPU devices, or both XPU and CUDA devices? I had to install mslk (on CUDA) so that the int4wo variants of the following Flux transformer TorchAO tests would not raise an mslk ImportError:

  • TestFluxTransformerTorchAo::test_torchao_quantization_num_parameters
  • TestFluxTransformerTorchAo::test_torchao_quantization_memory_footprint
  • TestFluxTransformerTorchAo::test_torchao_quantization_inference

class TestFluxTransformerTorchAo(FluxTransformerTesterConfig, TorchAoTesterMixin):
"""TorchAO quantization tests for Flux Transformer."""

modules_to_not_convert_for_test = ["proj_out"]

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.

Can you explain the motivation for not converting the proj_out layers here?

@dg845 dg845 left a comment

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.

Thanks! Left some comments.

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

Labels

size/S PR with diff < 50 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants