Reduce training failures from node failures and environment mismatches when training trillion-parameter models
When connecting multiple servers to train large language models, it is common for days of work to be lost due to a server stopping or library version mismatches. This tool bundles the entire training process, from GPU resource allocation to experiment deployment and checkpoint saving, into a single workflow. Developers can focus solely on model training code instead of complex infrastructure configuration.
Five roles in preventing training failures
This tool simultaneously acts as a GPU workload manager and a machine learning framework. It allocates node resources to users and supports ZeRO-3 and PyTorch's fully sharded data parallel API to help efficiently partition trillion-parameter models. It also provides a framework for starting, running, and monitoring large-scale neural network training, and manages resource conflicts through an experiment queue. Finally, it integrates with GitHub and GitHub Actions to support continuous integration for machine learning development.
Starting distributed training with a few lines of code
Training the LLaMa 70B model in a distributed environment requires only a few lines of code. Initialize the Llama70b model with zero_stage=3 and bf16 precision, and set up the AdamW optimizer and data loader. Inside a function decorated with @experiment, follow the standard PyTorch workflow of calculating the loss function in batches and updating the optimizer. Once training is complete, you can upload the model to the repository using the push_to_hub method.
Reducing the complexity of environments and configurations
After installing Docker, the project deployment key, and the higgsfield binary on the server, creating an experiment workflow ensures that code is deployed to nodes immediately upon being pushed to GitHub. The experiment execution UI is accessed via GitHub and automates the process of starting experiments and saving checkpoints. Because it follows the standard PyTorch workflow, it can be used alongside existing tools like deepspeed or accelerate, or applied to direct PyTorch sharding implementations. It avoids environment issues caused by different versions of PyTorch, NVIDIA drivers, or data processing libraries, and ensures reproducibility by documenting dependency versions and configurations. You can define experiments and design interaction methods through a simple interface without needing to define the 600+ arguments required for an experiment in a YAML file.
Node conditions to check before use
To use this tool, nodes must have the Ubuntu operating system installed and allow SSH access. A non-root user with passwordless sudo privileges is required. Testing has been completed on Azure, LambdaLabs, and FluidStack clouds, and issues can be opened for assistance if problems arise on other clouds. The command pip install higgsfield==0.0.3 is used for installation. Bug reports, feature requests, and installation or usage issues can be submitted via GitHub Issues and will receive a response within one day. New feature news is available on Twitter, while discussions and news can be found on the website.
