Train in epochs MLX-lm?

[mark-lord]

Hey all! Having come from training via Unsloth I was expecting epochs to be the standard, but running !python -m mlx_lm.lora --help I didn't immediately see a flag for training in epochs. GPT-4 helped me understand that the number of times the data will be seen is a function of the number of entries, the iters and number of batches. I hacked together some kinda inelegant code to make it possible to train based on epochs (or at least, so long as my understanding of all this is correct). But it's pretty ugly.

Was just wondering if there was some flag I was missing?

[mzbac]

mlx-lm lora may not have built-in support for that, but it could be achieved by calculating epochs based on iterations and batch size, adjusting them to achieve similar fine-tuning epochs.
for example

iters_per_epoch = len(train_set) // batch_size
iters = iters_per_epoch * epochs

[mark-lord]

That's a much much shorter way of doing it 😂 I'd been using this code:

def count_data_entries(file_path):
    """Count the number of entries in a JSON Lines file."""
    with open(file_path, 'r', encoding='utf-8') as file:
        return sum(1 for _ in file)  # Counts the number of lines

def calculate_iterations(data_size, batch_size, epochs):
    """Calculate the total number of iterations based on data size, batch size, and epochs."""
    iterations_per_epoch = -(-data_size // batch_size)  # Ceiling division for one epoch
    total_iterations = iterations_per_epoch * epochs
    return total_iterations

And then using that for setting iters. Which has been getting the job done, but I'll probably swap over to yours since I'm trying to make a beginner-friendly Jupyter Notebook and this will probably confuse people a lot more than yours 😂 That said, according to GPT-4, your code versus my code:

"In summary, with a batch size of 4 and 13 entries, mzbac's code snippet would result in 3 batches per epoch, ignoring the last remaining entry each time, while your code snippet would produce 4 batches per epoch, ensuring every entry is included, even if the last batch is smaller."

So if you're doing bigger batches and want to have everything seen in the epochs equally, your suggested code wouldn't cover that since you're doing ceiling division. Not sure if there's a better way to do this, but it could be changed to accomodate by doing:

iters_per_epoch = -(-len(train_set) // batch_size)
iters = iters_per_epoch * epochs

"This code ensures that even if the last batch is smaller than the other batches, it still counts as a separate iteration."

That said, from my own experiments I haven't seen any speed-ups from doubling up batch size, even if it all fits into memory. So keeping batch to 1 makes the most sense. In which case there's no difference between any of these ways of doing things.

[mark-lord]

Ohh weirdly I just implemented your code and now it's not properly reading the number of entries in the JSONL. It's meant to be 140 entries, but for some reason it's detecting 18 as iters_per_epoch? That's both with my version of your code + your original code. Going to try debug but clearly something is up with the way I implemented it. Unless I'm misunderstanding iters.

EDIT: figured out it's because as is, your code is counting the number of letters in the string that points to the file! Had to rejig it:

train_file = "./data/train.jsonl"
iters_per_epoch = -(-sum(1 for _ in open(train_file)) // batch_size)
iters = iters_per_epoch * epochs

This is finally working as intended 😊