Proposal: reuse methods in recipes

[Andrei-Aksionov]

Hi there 👋

Back in March, I left a comment under a PR that tries to refactor recipes, to make them simpler.

My idea (which is quite straightforward) is to avoid recipes that include complete code, essentially deviating from the single-file approach.
Let’s consider full-time fine-tuning and LoRA fine-tuning as examples. In this case, only the model instantiation and checkpoint saving should differ.
The proposal is to have full fine-tuning as the basis, and in the LoRA variant, only override the LoRA-specific methods.

The advantages and disadvantages, plus the example code, can be found in this PR.

---

Now, let’s discuss another matter.

If we adhere to this approach and:

1. Maintain full fine-tuning as the basis
2. Override only the LoRA-specific methods in the LoRA fine-tuning variant
3. Override only the distributed-specific methods in the full fine-tuning variant

… what should we do with the LoRA fine-tuning distributed variant?

• Option A: Override both the LoRA- and distributed-specific methods, which essentially means duplicating the entire code, and that brings us back to square one.
• Option B: Reuse the LoRA-specific code from the LoRA recipe and the distributed-specific code from the distributed recipe. In this case, to read the LoRA distributed recipe, we would need to jump between three files, which is too cumbersome and not ideal.

Another option would be to create hardware environment-agnostic code. Regardless of whether there is a single GPU or multiple GPUs, the code will remain the same. All the logic will be encapsulated in a separate class.

This approach would eliminate the need for separate recipes for distributed cases.

---

The question is what the core team thinks about it?

[pbontrager]

Hi @Andrei-Aksionov having single file implementations of our recipes is a key feature that we like to offer in order for people to be able to easily copy and edit the single file they're working on and not have to span and look through multiple files. That said, for ease of maintainability and to better keep features in sync across multiple recipes, it's good to revisit this from time to time. We plan to iterate on our recipe design later in the year, so please keep a look out for that and we'll definitely consider these proposals.

We want to consider whether it makes sense to keep separate distributed and single GPU recipes, where single GPU recipes primarily provide a less complex recipe for people to work with. Second we want to discuss how to best keep our recipes in sync. Some ideas here are using templates and linting style rules to enforce parts of recipes stay in sync with each other or generating recipes from parts when a user wants to copy or run them. Or we might decide that keeping things as they are is worth the extra effort for user simplicity.

[Andrei-Aksionov]

Yup, that's what I thought. Just wanted to know your opinion on that.

We plan to iterate on our recipe design later in the year, so please keep a look out for that and we'll definitely consider these proposals.

Sure.

---

Some ideas here are using templates and linting style rules to enforce parts of recipes stay in sync

One of the options might be to parse methods, that we know should be identical, and compare them in a test.

-- file_1.py

class MyClass:
    def method_different(self):
        return "This is a different method in MyClass in file_1.py"

    def method_same(self):
        return "This is a method in MyClass in file_1.py with the same name as in file_2.py"

-- file_2.py

class MyClass:
    def method_different(self):
        return "This is a different method in MyClass in file_2.py"

    def method_same(self):
        return "This is a method in MyClass in file_1.py with the same name as in file_2.py"

and the test will be something like this (but prettier, it's a quick draft).

-- test_recipe_...py

import ast
import inspect
import textwrap
import difflib
import pytest

from src.file_1 import MyClass as ClassA
from src.file_2 import MyClass as ClassB

def diff_strings(a: str, b: str, fromfile="File A", tofile="File B") -> str:
    """Generates a unified diff of two strings."""
    a_lines = a.strip().splitlines()
    b_lines = b.strip().splitlines()
    return "\n".join(
        difflib.unified_diff(
            a_lines, b_lines, fromfile=fromfile, tofile=tofile, lineterm=""
        )
    )


def get_method_ast(cls, method_name):
    """
    Extracts the AST for the specified method from a class.
    """
    # Get the method object from the class
    method = getattr(cls, method_name)
    # Get the source code of the method
    source = inspect.getsource(method)
    source = textwrap.dedent(source).strip()
    # Parse the source into an AST
    tree = ast.parse(source)

    # If the method is defined inside the class, the source might contain indentation.
    # The AST tree should contain a FunctionDef node in its body.
    for node in ast.iter_child_nodes(tree):
        if isinstance(node, ast.FunctionDef):
            return node
    raise ValueError(f"No method definition for {method_name} found in {cls.__name__}")


@pytest.mark.parametrize("method_name", ["method_same"])
def test_shared_method_ast_equal(method_name):
    ast1 = get_method_ast(ClassA, method_name)
    ast2 = get_method_ast(ClassB, method_name)

    unparsed_a = ast.unparse(ast1)
    unparsed_b = ast.unparse(ast2)

    if unparsed_a != unparsed_b:
        diff = diff_strings(
            unparsed_a, unparsed_b, fromfile="file_1.py", tofile="file_2.py"
        )
        pytest.fail(f"Method '{method_name}' differs:\n\n{diff}")