Skip to content

[RFC] Helion Remote Autotuning Service #2745

Description

@mrodden

During discussions with several folks about some of the challenges of autotuning, one of the repeated themes that came up was the challenge of being able to easily run Helion autotuning on production hardware due to availability of the hardware. Availability is mostly impacted by many newer platforms that are not widely available (e.g. B200), and often reserved for production environments which come with access restrictions to developers.

One solution that came up to solve this was the idea of a remote autotuning service that could run the autotuning asynchronously and remotely. This would enable developers to utilize multiple platforms for autotuning, allow for possible horizontal scaling of the autotuning process (like Mark Saroufim’s prototype), and possibly enable tooling like CI systems to run autotuning jobs in production during offpeak hours. This is just a handful of possible use cases that have come up from discussions.

A goal of this RFC is to start the discussion around such a service and gather/find the primary use cases from folks across the community.

A first prototype of the service that I have been considering creating as a “vertical slice” of the full concept is the following. The one use case I have during development of Helion kernels is being able to deploy Helion configs to different hardware (think sm100, sm120, sm121a, etc) and run them asynchronously from my kernel development environment, such that I am not tying up current the machine I am developing on, which I use to continue to work on other kernels locally, while the remote jobs run the search for optimal configurations. I think this would cover most of the primary implementation that would need to be created, while allowing extension and generalization with later work. The goal here would be to find the big design tradeoffs and also the parts where things like remote infrastructure can plug in to a general interface.

Primary desiderata I have currently:

  • Run autotuning jobs remotely on various hardware configurations and get the results of autotuning jobs back for capture and use
  • A way to deploy jobs to nodes that exist as static nodes, as well as dynamic nodes from an IaaS or PaaS provider (Modal, AWS EC2, Kubernetes/OpenShift, etc)
    • I think this is something that can be later added/extended by community contributions. Everyone seems to have their preferred platform.
  • Ability to scale job deployment across homogenous nodes (or close to homogenous) to speed up autotuning work (Mark’s prototype use case)
  • Possibly a priority queue system for jobs on restricted availability nodes. I.e. production B200 that could be a target for many developer autotuning jobs. Optionally, have the ability to only run on off peak times, or restrict the load on the production system.
  • Ability to upload results of autotuning jobs to a remote cache system for Helion kernels
  • Generalize the service to be able to work with different autotuning needs other than Helion (cutile, triton, etc)
    • This one is idealistic but I figured its worth calling out if someone has a real need

There are many possible ways to implement such a service, hence the focus on use cases, requirements, and other desiderata so far. Some implementation / design details that have come up so far:

A remote process with a gRPC frontend would be fairly easy to implement as a working prototype. Also a remote shell execution environment would be easy to implement and can work with pretty much anything with an SSH service, very similar to how Ansible works. gRPC or some HTTP interface would be a more formal interface and therefore more pluggable, but has security concerns with TLS and such.

There are some open questions on what the remote environment would need to look like:

  • Do we need a full Helion environment with pytorch, or do we just run the kernels?
  • How to store and retrieve results? Locally or remotely? Could be pluggable like with the remote cache
  • If we need a full environment, do we deploy with a container image or just use what is provided? This probably is going to depend on the underlying infrastructure provided. I.e. kubernetes clusters probably have a provided image to use already

Activity

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions