[CPU] Fix issue: No module named 'fbgemm_gpu.experimental'

[shiyang-weng]

We want to use torchrec on CPU. On CPU, torchrec dependent on the CPU version of fbgemm_gpu.
On CPU, fbgemm_gpu not include fbgemm_gpu.experimental.
Similar issue: #2591
env:

pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
pip3 install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cpu
pip3 install --pre fbgemm_gpu --index-url https://download.pytorch.org/whl/nightly/cpu

Reproduce: import torchao
Error:

  File "/home/wengshiy/ao/torchao/quantization/quantize_/workflows/int4/int4_preshuffled_tensor.py", line 25, in <module>
    if not _is_fbgemm_gpu_genai_available():
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/home/wengshiy/ao/torchao/utils.py", line 1162, in _is_fbgemm_gpu_genai_available
    import fbgemm_gpu.experimental.gen_ai  # noqa: F401
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'fbgemm_gpu.experimental'

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3292

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit f8e50ab with merge base 01374eb ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[shiyang-weng]

@Xia-Weiwen @mingfeima Could you help review this PR?

[jerryzh168]

oh maybe just do importlib.util.find_spec(fbgemm_gpu.experimental.gen_ai) is None?

[shiyang-weng]

Need to check fbgemm first, otherwise the following error will occur.

  File "/home/wengshiy/ao/torchao/quantization/quantize_/workflows/int4/int4_preshuffled_tensor.py", line 25, in <module>
    if not _is_fbgemm_gpu_genai_available():
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/home/wengshiy/ao/torchao/utils.py", line 1159, in _is_fbgemm_gpu_genai_available
    importlib.util.find_spec("fbgemm_gpu.experimental") is None
    ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib.util>", line 91, in find_spec
ModuleNotFoundError: No module named 'fbgemm_gpu'

[shiyang-weng]

@pytorchbot label "topic: not user facing"

[shiyang-weng]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: 1 mandatory check(s) are pending/not yet run. The first few are:

• Facebook CLA Check

Dig deeper by viewing the pending checks on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: superuser