Mac support

[wangkuiyi]

It looks like this project only releases packages for Linux, based on the files listed at PyPI. Am I correct? If so, is the project specifically designed for Linux, making a macOS release unnecessary? Thanks!

[d4l3k]

@wangkuiyi this will likely work just fine on Macs but it hasn't been a priority. Macs generally only have a single GPU and I'm not aware of anyone clustering them to do large scale distributed training.

We probably should provide a Mac release just to avoid dependency issues when developing on larger projects though

@wangkuiyi are you interested in helping out with this? I don't have a Mac so can't test this myself. We also would need to add a Mac build in our nightly package release CI

[wangkuiyi]

Hi Tristan,

Thank you for the explanation!

Yes, I’m interested in making TorchFT work on macOS if a macOS node can be added to the CI. I understand that macOS isn’t commonly used for distributed training, but enabling support could still be valuable—particularly for ensuring that VS Code and Cursor running on a Mac can properly resolve dependencies when parsing the source code of TorchTitan and other projects that depend on TorchFT.

I’ve been working on building TorchFT and running both the Python and Rust unit tests on macOS. So far, it mostly works, and it seems that the only minor modifications are needed to disable tests requiring NCCL.

Below are the steps I followed on macOS (M2 Ultra) for your reference. Please let me know if and how I should proceed.

Steps to Build and Test TorchFT on macOS

1. Install Rust

   Using the same shell script as in README.md. Need the following post-steps:

   . "$HOME/.cargo/env"  # For sh/bash/zsh/ash/dash/pdksh  
   source "$HOME/.cargo/env.fish"  # For fish  
   source "$HOME/.cargo/env.nu"  # For nushell  
   

2. Install Protobuf Compiler and Library

   brew install protobuf  

3. Install TorchFT

   pip install -e .[dev]

4. Run Python Tests

   Following TorchFT’s contribution guide, I ran:

   pytest torchft/

   However, the following tests failed due to the lack of NCCL on macOS:

   ERROR torchft/checkpointing/http_transport_test.py  
   ERROR torchft/checkpointing/pg_transport_test.py  
   ERROR torchft/checkpointing/rwlock_test.py  
   ERROR torchft/checkpointing/transport_test.py  
   ERROR torchft/coordination_test.py  
   ERROR torchft/data_test.py  
   ERROR torchft/ddp_test.py  
   ERROR torchft/device_mesh_test.py  
   ERROR torchft/fsdp_test.py  
   ERROR torchft/futures_test.py  
   ERROR torchft/lighthouse_test.py  
   ERROR torchft/local_sgd_integ_test.py  
   ERROR torchft/local_sgd_test.py  
   ERROR torchft/manager_integ_test.py  
   ERROR torchft/manager_test.py  
   ERROR torchft/multiprocessing_test.py  
   ERROR torchft/optim_test.py  
   ERROR torchft/parameter_server_test.py  
   ERROR torchft/process_group_test.py  
   

   All failures were caused by:

   ImportError: cannot import name 'ProcessGroupNCCL' from 'torch.distributed'  
   (/opt/homebrew/Caskroom/miniforge/base/envs/torchft/lib/python3.10/site-packages/torch/distributed/__init__.py)  
   

5. Build Rust Code and Run Rust Tests

   I ran the following command:

   PYTHON_SYS_EXECUTABLE="/opt/homebrew/Caskroom/miniforge/base/envs/torchft/bin/python" \  
   PYO3_PYTHON="$PYTHON_SYS_EXECUTABLE" \  
   RUSTFLAGS="-C link-arg=-undefined -C link-arg=dynamic_lookup" \  
   cargo test  

   Note: When working with miniforge, these environment variables are required; otherwise, cargo won’t be able to locate PyO3, which is necessary for building the Python bindings.

⸻

Please let me know how you’d like to proceed or if any additional modifications are needed.

Best,
Yi

[d4l3k]

@wangkuiyi Thanks for digging into this!

It should be pretty simple to make ProcessGroupNCCL a conditional import. We may want to move all the NCCL code into a new process_group_nccl.py file for simplicity

Also good to know on the Rust side of things, should be simple enough to add those for builds

I think it would be a good first step to set this up on the github CI for the unittest target. We have access to mac GitHub Actions runners so we can update the build matrix to include a build on those. Ex: https://github.com/pytorch/torchx/blob/main/.github/workflows/python-unittests.yaml#L15-L17

It may be annoying for you to try and get this working since I'll have to manually run the CI for you. If you put up a PR with the Mac fixes I can get the CI build setup

[wangkuiyi]

I’m happy to file a PR. However, it doesn’t seem like the solution is as simple as collecting and optionally excluding NCCL-related tests. All the failed tests encountered the same error:

______________________________________________________________________________ ERROR collecting torchft/process_group_test.py ______________________________________________________________________________
ImportError while importing test module '/Users/y/w/torchft/torchft/process_group_test.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
/opt/homebrew/Caskroom/miniforge/base/envs/torchft/lib/python3.10/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
torchft/__init__.py:8: in <module>
    from torchft.ddp import DistributedDataParallel
torchft/ddp.py:26: in <module>
    from torchft.process_group import ProcessGroup, ProcessGroupDummy, ProcessGroupGloo
torchft/process_group.py:45: in <module>
    from torch.distributed import (
E   ImportError: cannot import name 'ProcessGroupNCCL' from 'torch.distributed' (/opt/homebrew/Caskroom/miniforge/base/envs/torchft/lib/python3.10/site-packages/torch/distributed/__init__.py)

This is due to process_group.py importing ProcessGroupNCCL from torch.distributed.

torchft/torchft/process_group.py

Lines 45 to 50 in 9d38340

	from torch.distributed import (
	DeviceMesh,
	PrefixStore,
	ProcessGroup as BaseProcessGroup,
	ProcessGroupGloo as BaseProcessGroupGloo,
	ProcessGroupNCCL as BaseProcessGroupNCCL,

[d4l3k]

@wangkuiyi you'll need to make the tests AND imports conditional i.e. moving the ProcessGroupNCCL/ProcessGroupBabyNCCL to a separate file for instance

It's the same way we handle this in core PyTorch Distributed but it is a bit annoying to restructure things

[wangkuiyi]

I’m currently using a quick (and probably hacky) solution to allow me to make some progress and identify any potential issues.

I wrapped the import of ProcessGroupNCCL in process_group.py with the following condition (main...wangkuiyi:torchft:make_nccl_optional):

if dist.is_nccl_available():
    from torch.distributed import ProcessGroupNCCL as BaseProcessGroupNCCL

This workaround helped me uncover a bug in PyTorch Distributed: pytorch/pytorch#149537

I have also submitted a fix for it: pytorch/pytorch#149538.

Could you please take a look? Thanks!