FMS Acceleration for Fused Operations and Kernels

This library contains fused operations and custom kernels, to be expanded over time. Currently it contains the following:

1. Fused operations and kernels extracted from unsloth.
   • Low-Rank Adapter Fused Operations
   • Fast RoPE Triton Kernels
   • Fast RMS LayerNorm Triton Kernels
   • Fast Cross Entropy Triton Kernels

Plugins

Supported DataType Settings

Compatibility Matrix with Mixed Precision

NOTE: this chart is also a good reference for supported types, even for the non-FOAK case.

Code Extracted from Unsloth

Notes on the extraction of code from unsloth:

• While unsloth is released under Apache 2.0, there are comments indicating some exceptions strewn throughout the code base, see an example here.

  it would require a commercial license if used to run on more than 4 GPUs ...
  

• These exceptions appear to be located around the trainer improvements, see another example here.
• These exceptions appear around Feb 2024 Release; any code that appears in any file where such exceptions occur is not extracted.
• Instead in its place, we have adopted a different approach; we adopt the approach of model patching, as opposed unsloths' approach to rewrite the model. Our approach is novel and completely rewritten from scratch.
• We have also enabled dropout on the lora fused operations.
• All extracted code appears before the Feb 2024 Release.
• In the table below we record what was extracted, and the exact commit from which it was taken.

Supported Models

Adding Support For A New Model

It is realtively easy by following an existing template, in what follows we use GraniteForCausalLM as an example.

• implement a get_mp_rules for the new model, which returns a list of ModelPatcherRule.
• logic that needs to be changed is the various classes that the rules are triggered on. Import the various module classes likes so:

  from transformers.models.granite.modeling_granite import ( 
      GraniteAttention,
      GraniteMLP,
      GraniteRMSNorm,
  )
  

• replace the classes appropriately in various locations in ModelPatcherRule. In particular the ModelPatcherTrigger portions of it. Name rule_id appropriately.

  ModelPatcherRule(
      rule_id="granite-rms",
      trigger=ModelPatcherTrigger(check=GraniteRMSNorm),
      forward=fast_rms_layernorm,
  )
  

Known Issues

• MixedPrecision --fp16 or --bf16 should be used with fast_lora.
• fast_lora has issues with FSDP V1 with the peft style of FSDP wrapping.
  • This is because the adapter's forward functions are bypassed in the fused ops.
  • For AutoGPTQ/QLoRA this is addressed by distributing the adapters using DDP so they will be unsharded in time for the fused ops.
• fast_rope_embeddings does not work with position_ids. Currently position_ids are ignored and could give wrong results.