Add module connector for LoRA
State: closed | Pull request: yes
Labels: name: stale

# What does this PR do ?

The LoRA implementation is currently hard-coded to `TransformerEngine` and torch nn.linear layers. It fails on regular `RowParallelLinear` and `ColumnParallelLinear` because the attributes it looks for do not exist on those modules. This module connector gives a mechanism to override the get/set behavior of the LoRA transform on the torch modules it transforms.

While we can update the default connector behavior, this approach allows anyone who might want to integrate LoRA with new layer types to do so without having to override and copy/paste/edit the full transform method.

This MR also adds in support for mcore parallel linear layers, and removes the need to have Transformer Engine installed to use LoRA.

**Collection**: llm

# Changelog 
- Module connector added for LoRA transform.
- MCore layers supported for LoRA
- Transformer Engine installation not required for LoRA

# Usage

Here is an example of overriding the module connector with a connector that can choose the in_features based on which attributed the `module` has.

```python
class GeneFormerModuleConnector(LoRAModuleConnector):
    @staticmethod
    def get_in_features(m: nn.Module):
        match m:
            case object(input_size=n):
                return n
            case object(in_features=n):
                return n
            case _:
                raise ValueError(f"Module {m} does not have a supported input size calculation.")

    @staticmethod
    def get_out_features(m: nn.Module):
        match m:
            case object(output_size=n):
                return n
            case object(out_features=n):
                return n
            case _:
                raise ValueError(f"Module {m} does not have a supported output size calculation.")


class LoRAForGeneFormerTokenRegressor(LoRA):
    def __post_init__(self):
        self.module_connector = GeneFormerModuleConnector()

```

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


@ashors1 @marcromeyn any chance either of you could review or suggest an alternate reviewer?

It feels a bit arbitrary to me to have an extra abstraction just to get the number of input and output features. I would prefer if we refactor the transform method a bit into some more methods to get things like input and output features so people could subclass and overwrite those if needed. 

The current LoRA class is kinda GPT specific so it would make sense to me if there’s one for GeneFormer that subclasses this one and overwrite some methods.

> It feels a bit arbitrary to me to have an extra abstraction just to get the number of input and output features. I would prefer if we refactor the transform method a bit into some more methods to get things like input and output features so people could subclass and overwrite those if needed.
> 
> The current LoRA class is kinda GPT specific so it would make sense to me if there’s one for GeneFormer that subclasses this one and overwrite some methods.

when reading the code, the type signature on LoRA transform method makes it seem like there is no concept of input/output features -- nn.Module provides no such guarantee. The code as currently written implies there is a stateful transformation that _must_ occur (somewhere, and at the right time) to the module to set the `in_features` and `out_features` attributes. 

I find it much more intuitive when we can use type annotations to communicate expectations. In the case of PEFT, knowing the input/output shape is a pre-requisite. The abstraction @georgea provides makes this expectation explicit and provides sensible defaults for RowParallelLinear and ColumnParallelLinear. In my opinion the ideal solution is to use a new type for this, but im guessing this is difficult with Megatron and TE?


> > It feels a bit arbitrary to me to have an extra abstraction just to get the number of input and output features. I would prefer if we refactor the transform method a bit into some more methods to get things like input and output features so people could subclass and overwrite those if needed.
> > The current LoRA class is kinda GPT specific so it would make sense to me if there’s one for GeneFormer that subclasses this one and overwrite some methods.
> 
> when reading the code, the type signature on LoRA transform method makes it seem like there is no concept of input/output features -- nn.Module provides no such guarantee. The code as currently written implies there is a stateful transformation that _must_ occur (somewhere, and at the right time) to the module to set the `in_features` and `out_features` attributes.
> 
> I find it much more intuitive when we can use type annotations to communicate expectations. In the case of PEFT, knowing the input/output shape is a pre-requisite. The abstraction @georgea provides makes this expectation explicit and provides sensible defaults for RowParallelLinear and ColumnParallelLinear. In my opinion the ideal solution is to use a new type for this, but im guessing this is difficult with Megatron and TE?

Exactly. Basically, right now LoRA.transform(module: nn.Module), doesn't communicate what info is needed from a module to make the right lora transformation. What I've done is effectively mapped the problem onto LoRa.transform(module_connector: ModuleConnector) which is typed and composable. Such that, a developer knows what information is needed to make the transform on any module. Making it a separate abstraction from Lora isolates the module dependency and allows re-use of the LoRA specific logic.

This is very compatible with the idea you presented @marcromeyn , where the LoRA adaptation might need to be updated based on the specifics of the model. E.g., between GPT and T5 or something like that.

What this does is isolate the module specific logic from the actual implementation of the LoRA adapter.

In the version you propose, assume a developer needs to implement support for a new module type.

If you have a bunch of code like:

```python
class LoRAForGPTOnTE(LoRA):

  def transform(m: nn.Module, ...):
    if m.name in self.gpt_modules_to_adapt:
        return Adapter(self.get_out_features(m))

  def get_out_features(m):
    return m.out_features * tp_size

class LoRAForGPTOnMcore(LoRA):
  def transform(m: nn.Module, ...):
    if m.name in self.gpt_modules_to_adapt:
        return Adapter(self.get_out_features(m))

  def get_out_features(m):
    return m.output_size 

class LoRAForT5OnTE(LoRA):
  def transform(m: nn.Module, ...):
    if m.name in self.t5_modules_to_adapt:
      return Adapter(self.get_out_features(m))

  def get_out_features(m):
    return m.out_features * tp_size

class LoRAForT5OnMcore(LoRA):
  def transform(m: nn.Module, ...):
    if m.name in self.t5_modules_to_adapt:
      return Adapter(self.get_out_features(m))  

  def get_output_size(m):
    return m.output_size
```

The developer now has to go and make a new subclass for each model to support the backend.

If instead we write the LoRA transform to a class for the backend connector that standardizes the information LoRA needs to infer it's adapter, the update becomes trivial--just tell me how to get the info for your module.

So the logic becomes more like:

```python
class LoRAForGPT(LoRA):

  def transform(m: nn.Module, ...):
    if m.name in self.gpt_modules_to_adapt:
        return Adapter(self.module_connector.get_out_features(m))


class LoRAForT5(LoRA):
  def transform(m: nn.Module, ...):
    if m.name in self.t5_modules_to_adapt:
      return Adapter(self.module_connector.get_out_features(m))  

```

By writing just one class to connect to your module types, you get access to all pefts for all models by just overriding their `module_connector`. This is way less duplicated code, is more maintainable and extensible.

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.