Adding LLM support with peft +quantization (4bit and 8 bit)

[rajithkrishnegowda]

Summary

Adding LLM support with peft + quantization(4bit and 8bit) using several open source frameworks to improve local training (PEFT) by referencing FlowerLLM methods.

Type of Change (Mandatory)

Specify the type of change being made.

• Feature enhancement

Description (Mandatory)

This PR adds a new Jupyter notebook that demonstrates how to fine-tune Microsoft's Phi-4 model in a federated learning workflow using OpenFL. The notebook implements and compares both 4-bit and 8-bit quantization techniques along with Parameter-Efficient Fine-Tuning (PEFT) methods.

Key features of this notebook:

• Implements federated fine-tuning of Phi-4 using OpenFL's experimental workflow interface
• Applies QLoRA (Quantized Low-Rank Adaptation) for memory-efficient training
• Provides comparative analysis between 4-bit and 8-bit quantization in terms of:
• Memory usage across training phases
• Training and evaluation loss metrics
• Peak memory utilization
• Includes visualization tools to analyze performance differences
• Demonstrates gradient checkpointing and other memory optimization techniques
• Uses the math_10k.json dataset for fine-tuning

This notebook is intended for the experimental workflow section of OpenFL tutorials, helping users understand the trade-offs between different quantization approaches when fine-tuning large language models in federated settings.

Testing

The notebook has been tested on a system with CUDA support. All cells execute successfully, and both 4-bit and 8-bit quantization workflows complete the federated training process. The notebook generates visualizations comparing memory usage and performance metrics between the two quantization approaches.

https://github.com/rajithkrishnegowda/openfl/blob/rajith-llm-quanti/openfl-tutorials/experimental/workflow/LLM/phi-4-peft-quantization.ipynb

[rahulga1]

@rajithkrishnegowda Great Work. Can you also put metrics and plots for the baseline model, I dont think I see it currently. It will be really good to see the comparision.

[rajithkrishnegowda]

@rajithkrishnegowda Great Work. Can you also put metrics and plots for the baseline model, I dont think I see it currently. It will be really good to see the comparison.

added in latest commit

[kminhta]

This is great, thanks @rajithkrishnegowda ! Left some comments (mostly general questions and some minor comments) on my first walk through.

[kminhta]

Question: did you find that re-instantiating the trainer each round also reset the optimizer state? In Torch, you can checkpoint the optimizer and reload it with the model, but you usually need to explicitly do this. I'm not sure if the peft or trl libraries automatically handle this.

I don't think it's a major concern at this stage, but it's something to look into since it's not ideal to reset each time.

[rajithkrishnegowda]

Yeah good catch.

we can improve the implementation:

1. Extract and save the optimizer state during federated averaging
2. Restore the optimizer state when setting up the trainer in each round
3. Use PyTorch's built-in optimizer state dictionary methods

WDYT?

[kminhta]

Sounds like a good plan to me, thanks

[rajithkrishnegowda]

resolved in 46dba98

[kminhta]

This line seems misleading since 8-bit had a lower raw loss. Since it's a comment on efficiency, maybe just "4-bit provides more efficiency memory usage relative to loss"?

[rajithkrishnegowda]

yeah noted

[rajithkrishnegowda]

resolved in 46dba98

[kminhta]

Is it possible to plot a longer experiment or would that be too resource intensive? It's hard to capture the true trend in 2 rounds

Edit: actually never mind, it seems the general recommendation is 1-3 epochs for fine-tuning (https://docs.unsloth.ai/get-started/fine-tuning-guide#epochs) and since each round is a full epoch, it may not make sense to run too many more rounds (though I guess this also assumes optimal hyperparameters, data, etc. and federated LLMs might have different "best practices" due to inherent heterogeneity, but idk tbh). In that case, have you explored possibility of more frequent updates and if the trainer could support something like that? This may just be a generally open question, so no problem if you haven't

[rajithkrishnegowda]

Possible Approaches for More Frequent Updates

1. Partial Epoch Updates
   We could modify the train() method in the FederatedFlow class to support partial epoch training. Instead of calling trainer.train() which processes the entire dataset

2. Customizing SFTTrainer
   The TRL library's SFTTrainer could be extended to support this federated scenario

3. Modifying Batch Size/Data Sharding
   Another approach is to keep the one-epoch-per-round model but adjust how much data each collaborator processes in that "epoch"

We could implement an approach that allows configurable update frequency:

1. Add a parameter to the [FederatedFlow]class to control update frequency
2. Modify the train() method to process partial batches based on this frequency
3. Ensure optimizer state preservation between updates (already implemented)

WDYT?

[kminhta]

This seems like it could work. An alternative approach could be to modify max_steps in the TrainingArguments https://huggingface.co/docs/transformers/v4.51.3/en/main_classes/trainer#transformers.TrainingArguments.max_steps instead of epochs.

[rajithkrishnegowda]

with these change i.e max_steps, for one collaborator to complete one round for training takes one hour, so for 2 collaborator2 for 2 rounds for each quanitization takes 4 hours. Is this ideal?

in the existing approach, it takes 220 steps for training where as for max_steps it takes 1116 steps hence taking too much time.

[kminhta]

From offline discussion > set the max_steps to a value lower than the 220 steps it takes to run one epoch

[rajithkrishnegowda]

resolved in 46dba98

[MasterSkepticista]

General comment on the training output cells: There is too much detail in these cells which, as a reader, doesn't tell me why and what is the goal of this notebook. As noted in previous reviews, notebooks are a hook for a user to gain familiarity, objective and outcomes of the framework without running any code.

Therefore, an ideal notebook:

• Has to be short in length. Although this can be waived when tutorial itself covers many concepts.
• Defines objectives early on (this is taken care of in this notebook).
• Each cell must not exceed one scroll length worth of content, i.e. 15-20 lines - it fits the context window of the reader. If it does, consider breaking the cell.
• Output must be brief - reader knows what to expect in the output cell (as if they are executing the notebook in their mind)
• Conclusion must demonstrate that objectives defined at the beginning are met, through small, easy to parse plots.

Specific suggestions:

1. Please shrink/reformat/reconsider how much output text is presented.
2. Please run the experiment for at least 10 rounds to demonstrate a trendline. 2 rounds is not sufficient.

[rajithkrishnegowda]

General comment on the training output cells: There is too much detail in these cells which, as a reader, doesn't tell me why and what is the goal of this notebook. As noted in previous reviews, notebooks are a hook for a user to gain familiarity, objective and outcomes of the framework without running any code.

Therefore, an ideal notebook:

• Has to be short in length. Although this can be waived when tutorial itself covers many concepts.
• Defines objectives early on (this is taken care of in this notebook).
• Each cell must not exceed one scroll length worth of content, i.e. 15-20 lines - it fits the context window of the reader. If it does, consider breaking the cell.
• Output must be brief - reader knows what to expect in the output cell (as if they are executing the notebook in their mind)
• Conclusion must demonstrate that objectives defined at the beginning are met, through small, easy to parse plots.

Specific suggestions:

1. Please shrink/reformat/reconsider how much output text is presented.
2. Please run the experiment for at least 10 rounds to demonstrate a trendline. 2 rounds is not sufficient.

1. separated cell for 4bit and 8 bit but output of those cells can't be restricted in the settings of visual code when we commit it can't be restricted. Please let me know if you have any ideas
2. As per kevin's comment above since each round is a full epoch, it may not make sense to run too many more rounds but still ran 5 rounds.

[MasterSkepticista]

Plotting is a net improvement, however, notebook is still huge (12k lines of source currently!). The output cell of flflow alone takes 90% of the entire notebook length

Suggestions in no particular order:

• Replace GPU driver installation steps with a generic URL. I wouldn't mind removing it completely either...
• Merge pip installation cells and use cell comments if required
• Remove try-except blocks in the main code. For example, we need not try to load the optimizer state dictionary. If it fails, exception is raised by torch anyway. This will save lines used in this notebook's exception handling. Another relevant example:

  # Simplify training for 8-bit
  try:
      # Use trainer.train() for simpler training flow
      trainer.train(resume_from_checkpoint=False)
      # Get the last loss
      self.loss = trainer.state.log_history[-1].get('loss', float('inf'))
  except Exception as e:
      print(f"Training failed with error: {str(e)}")
      self.loss = float('inf')  # Set to infinity to indicate failure

• Migrate discrete sections to a new trainer_utils.py. One example would be loss calculation:

  def causal_crossentropy_loss(logits, labels):
      shift_logits = logits[..., :-1, :].contiguous()
      shift_labels = labels[..., 1:].contiguous()
      loss = F.cross_entropy(
          shift_logits.view(-1, shift_logits.size(-1)),
          shift_labels.view(-1),
          ignore_index=-100
      )
      return loss
  # Use fn call in the notebook...

• Consider shrinking log/print messages. It requires horizontal scroll on some:

  print(f"[Round {self.current_round}, Update {self.current_sub_round}] Average local model validation loss = {self.local_model_accuracy}")

• kB/MB typo in the plots
• Loss must be decreasing each round, one of the plot shows increasing loss trend
• 8-bit mode training loss missing in last two plots
• Evaluation loss trend of 4-bit is lower than 8-bit: I think it should be opposite (more bits, better generalization?)

[MasterSkepticista]

Nit: rename as plotting_utils.py or utils.py

[teoparvanov]

As pointed out by @MasterSkepticista, it looks like GitHub is struggling to display the diff for the new notebook. This is most in part due to the bloat related to stdouts. Please exclude those as a first step. Next, you could consider extracting more of the utility methods to python modules outside of the notebook.

[teoparvanov]

Additionally, as discussed offline, I'd suggest excluding self.model between the FL flow steps and only transmitting the PEFT/LoRA parameters instead. In fact, the rest of the model is anyway supposed to be "frozen" during fine-tuning, so one more reason not to transmit it entirely back and forth between the OpenFL nodes.

This will reduce unnecessary large object cloning within the LocalRuntime, but more importantly - it will enable running the notebook with FederatedRuntime, and also via TaskRunner API where we have a 2GB limitation on the protobuf message size.

[porteratzo]

Great work on the PR! I have a few suggestions:

SFTTrainer in Cell 18: It seems like the SFTTrainer might not be necessary. Could you check if we can remove it to simplify the code?

Unused Imports: I noticed some unused imports in cells 8 and 23. Cleaning these up could improve the code's readability.

Optimizer Initialization: There are two sections for initializing the optimizer. Is the second one required, or can we consolidate them to avoid redundancy?

Memory Optimization: Since you're using localruntime, you might reduce memory usage by sending the model to the CPU before starting the collaborator steps. Loading it into the GPU for these steps and then sending it back to the CPU for the join step could be more efficient. Without that you are basically holding 3 models in memory at the same time, (1 for aggregator and 2 for collaborators)

Let me know what you think!