Cleanup setup.py

[eitanturok]

This PR improves setup.py in several ways:

1. It fixes a small type error with extra_deps['all'].
2. It adds torch as a requirement in install_requires. This was missing before even though the library requires torch.
3. We deleted the import of torch in setup.py. Since setup.py does not actually use any torch functions, this import is unnecessary.

[tgale96]

Hey, looks awesome! Thanks for the PR! Apologies for the delay on this an others some access perms seems to be messed up currently...

Give me the thumbs up and I'll merge!

[tgale96]

I heard you hit an issue needing torch at install time? Could you share the error?

[eitanturok]

Hi @tgale96.

Just saw your message. Thanks so much for getting back to me.

This PR is ready and I'd love it if you could merge it.

Background:
The error I got was the result of some complexities introduced by adding a pyproject.toml to MegaBlocks. This changes how pip installs packages in isolation (more info here 1, 2) which causes the computer to think that torch is not installed even when it is.

TLDR: After adding a pyproject.toml to MegaBlocks, running

pip install "torch<2.4"
pip install .

in the MegaBlocks repo in a fresh virtual environment failed with the error

(megablocks-16) root@62d6f7c1-f73e-4060-96e7-3fb0da4091cd-0:/mnt/workdisk/eitan/10_productionize_megablocks/megablocks# pip install . --no-cache-dir  
Processing /mnt/workdisk/eitan/10_productionize_megablocks/megablocks
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Installing backend dependencies ... done
  Preparing metadata (pyproject.toml) ... done
Collecting numpy<2.1.0,>=1.21.5 (from megablocks==0.5.1)
  Obtaining dependency information for numpy<2.1.0,>=1.21.5 from https://files.pythonhosted.org/packages/ef/27/39622993e8688a1f05898a3c3b2836b856f79c06637ebd4b71cb35cc9b18/numpy-2.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata
  Downloading numpy-2.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (60 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 60.9/60.9 kB 162.1 MB/s eta 0:00:00
Requirement already satisfied: torch<2.4,>=2.3.0 in /mnt/workdisk/eitan/10_productionize_megablocks/venvs/megablocks-16/lib/python3.11/site-packages (from megablocks==0.5.1) (2.3.1)
Requirement already satisfied: triton>=2.1.0 in /mnt/workdisk/eitan/10_productionize_megablocks/venvs/megablocks-16/lib/python3.11/site-packages (from megablocks==0.5.1) (2.3.1)
Collecting stanford-stk==0.7.0 (from megablocks==0.5.1)
  Downloading stanford-stk-0.7.0.tar.gz (17 kB)
  Installing build dependencies ... done
  Getting requirements to build wheel ... error
  error: subprocess-exited-with-error
  
  × Getting requirements to build wheel did not run successfully.
  │ exit code: 1
  ╰─> [20 lines of output]
      Traceback (most recent call last):
        File "/mnt/workdisk/eitan/10_productionize_megablocks/venvs/megablocks-16/lib/python3.11/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 353, in <module>
          main()
        File "/mnt/workdisk/eitan/10_productionize_megablocks/venvs/megablocks-16/lib/python3.11/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 335, in main
          json_out['return_val'] = hook(**hook_input['kwargs'])
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/mnt/workdisk/eitan/10_productionize_megablocks/venvs/megablocks-16/lib/python3.11/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 118, in get_requires_for_build_wheel
          return hook(config_settings)
                 ^^^^^^^^^^^^^^^^^^^^^
        File "/tmp/pip-build-env-j55yqtzx/overlay/lib/python3.11/site-packages/setuptools/build_meta.py", line 327, in get_requires_for_build_wheel
          return self._get_build_requires(config_settings, requirements=[])
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/tmp/pip-build-env-j55yqtzx/overlay/lib/python3.11/site-packages/setuptools/build_meta.py", line 297, in _get_build_requires
          self.run_setup()
        File "/tmp/pip-build-env-j55yqtzx/overlay/lib/python3.11/site-packages/setuptools/build_meta.py", line 497, in run_setup
          super().run_setup(setup_script=setup_script)
        File "/tmp/pip-build-env-j55yqtzx/overlay/lib/python3.11/site-packages/setuptools/build_meta.py", line 313, in run_setup
          exec(code, locals())
        File "<string>", line 2, in <module>
      ModuleNotFoundError: No module named 'torch'
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
error: subprocess-exited-with-error

× Getting requirements to build wheel did not run successfully.
│ exit code: 1
╰─> See above for output.

note: This error originates from a subprocess, and is likely not a problem with pip.

Removing torch from the setup.py of stk fixes this issue.

Thanks for being so responsive. Cheers!

[eitanturok]

There exists a similar problem in grouped_gemm and I made a similar PR to fix it.

I'd appreciate if you could check out that PR too.

Thanks!

[tgale96]

Thanks, Eitan! I appreciate the detailed explanation!