Show HN: Run PyTorch locally with a remote GPU backend

5 points by alyxya a day ago

I integrated a remote GPU execution backend into PyTorch through the same system that custom hardware accelerators get integrated into PyTorch. You can create a remote machine and obtain its CUDA device whenever you want to create or move tensors onto the remote GPU.

  machine = mycelya_torch.RemoteMachine("modal", "A100")
  cuda_device = machine.device("cuda")
  x = torch.randn(1000, 1000, device=cuda_device)
  y = torch.randn(1000, 1000).to(cuda_device)

  import mycelya_torch
  from transformers import AutoModelForCausalLM, AutoTokenizer

  @mycelya_torch.remote
  def load_model(model_name: str):
      tokenizer = AutoTokenizer.from_pretrained(model_name)
      model = AutoModelForCausalLM.from_pretrained(
          model_name, torch_dtype="auto", device_map="auto"
      )
      return model, tokenizer