Examples

Training

Fine-tuning

This example shows how to fine-tune a causal langauge model using AxoNN as a backend to accelerate.

To use AxoNN with the Hugging Face Accelerate API, you need to declare the AxoNN plugin and specify the required arguments:

from accelerate import AxoNNPlugin
axonn_plugin = AxoNNPlugin(G_intra_depth=args.tensor_parallelism, G_intra_col=1, G_intra_row=1)
accelerator = Accelerator(..., axonn_plugin=axonn_plugin, ...)

In the above code block, G_intra_depth is, G_intra_col is, and G_intra_row is

In addition, our context manager is used to parallelize linear layers within the model architecture. Essentially, this is performing the monkey patching mentioned in the EasyAPI section within the User Guide automatically. This enables AxoNN to fine-tune the given model in parallel.

with axonn.models.transformers.parallelize(args.model_name_or_path):
        model = AutoModelForCausalLM.from_pretrained(
            args.model_name_or_path,
            from_tf=bool(".ckpt" in args.model_name_or_path),
            config=config,
            low_cpu_mem_usage=args.low_cpu_mem_usage,
            trust_remote_code=args.trust_remote_code,
        )

Inference