Fine-Tuning Federado de Modelos de Linguagem

Esse laboratório apresenta o Fine-tuning federado de modelo de linguagem para uma aplicação de detecção de anomalias em Logs de sistemas. O laboratório utiliza como base o trabalho Fine-Tuning Eficiente de Modelos de Linguagem Para Detectar Anomalias em Logs Privados usando Aprendizado Federado e também baseado no conteúdo do minicurso MC2 - Fine-tuning de LLMs aprensetado no SBRC 2025

Para executar o laboratório utilize o Google Colab

Dependências

from google.colab import drive
drive.mount('/content/drive/', force_remount=True)

import sys

sys.path.append('/content/drive/MyDrive/Colab Notebooks/Minicurso SBRC')  # Path to your folder

Instalação de Bibliotecas

!pip install -q transformers sentencepiece evaluate "flwr[simulation]" trl bitsandbytes

!export WANDB_DISABLED=true

Download e preparação do dataset

• Crie uma pasta dataset no seu drive.
• Copie o conteúdo do diretório compartilhado para a pasta criada:
  • Link para o drive

Imports

from flwr.common import Context
from flwr.server import ServerConfig
import torch
import os
import json
from datasets import Dataset
from flwr.client import NumPyClient, ClientApp
from flwr.common import ndarrays_to_parameters
from flwr.server import ServerAppComponents, ServerApp
from flwr.server.strategy import FedAvg
from flwr.simulation import run_simulation

from utils import FitConfigFactory, get_evaluate_fn
from utils import set_parameters, cosine_learning_rate, TraningConfigBuilder, TrainerBuilder, get_parameters, \
    ModelBuilder, get_tokenizer

Variáveis de configuração da simulação

#paths
testset_path = "/content/drive/MyDrive/Colab Notebooks/Minicurso SBRC/dataset/test.csv"
results_path = "/content/drive/MyDrive/Colab Notebooks/Minicurso SBRC/results"
dataset_path = "/content/drive/MyDrive/Colab Notebooks/Minicurso SBRC/dataset"

#model
model_name = "HuggingFaceTB/SmolLM-135M"
lora_rank = 8
lora = True
initial_lr = 1e-3
min_lr = 1e-5

#training
num_supernodes = 5
num_rounds = 2
fraction_fit = 0.4
fraction_eval = 0.0

train_context_dict = {
    "num-rounds": num_rounds,
    "initial-lr": initial_lr,
    "min-lr": min_lr,
    "dataset-path": dataset_path,
    "results-path": results_path,
    "model-name": model_name,
    "lora": lora
}

#eval
experiment_name = "experimento_sbrc"

eval_context_dict = {
  "model-name": model_name,
  "lora-rank": lora_rank,
  "lora": lora,
  "testset-path": testset_path,
  "results-path": results_path,
  "nrows": None,
  "experiment-name": experiment_name
}

if torch.cuda.is_available():
  device = torch.device("cuda")
  os.environ["CUDA_VISIBLE_DEVICES"] = "0"
  torch.cuda.empty_cache()
else:
  device = torch.device("cpu")

Cliente

class LLMClient(NumPyClient):
    def __init__(self, cid, model, tokenizer) -> None:
        super().__init__()

        self.cid = cid
        self.model = model
        self.tokenizer = tokenizer

    def fit(self, parameters, config):
        current_round = config["current_round"]
        total_rounds = config["num_rounds"]
        initial_lr = config["initial_lr"]
        min_lr = config["min_lr"]
        dataset_path = config["dataset_path"]
        results_path = config["results_path"]
        model_name = config["model_name"].lower()
        sim_name = "experimento_sbrc"
        lora = config["lora"]

        # Obtém dataset do cliente
        client_dataset = Dataset.load_from_disk(f"{dataset_path}/client_{self.cid}")
        # Obtém o modelo a ser treinado
        set_parameters(self.model, parameters, lora)
        self.model.to(device)

        # Configuração do treinamento
        # Calcula a nova taxa de aprendizado
        new_lr = cosine_learning_rate(current_round=current_round,
                                      total_rounds=total_rounds,
                                      initial_lr=initial_lr,
                                      min_lr=min_lr)
        # Cria a configuração do treinamento
        training_args = TraningConfigBuilder().with_output_dir(results_path).with_logging_dir(
            results_path).with_learning_rate(new_lr).build()

        # Cria objeto responsável por treinar o modelo
        trainer = TrainerBuilder().with_cid(self.cid).with_model(self.model).with_args(
            training_args).with_train_dataset(client_dataset).with_tokenizer(
            self.tokenizer).with_eval_dataset(client_dataset).with_model_name(model_name).build()

        # Realiza treinamento
        print(f"Rodada {current_round}: Treinando Cliente {self.cid} com lr {new_lr}")

        trainer.train()

        parameters = get_parameters(self.model, lora)
        dataset_size = len(client_dataset['labels'])

        print(f"Rodada {current_round}: Cliente {self.cid} treinou")

        # Save losses
        output_dir = f"{results_path}/fl-results/{sim_name}/round_{current_round}/client_{self.cid}"
        os.makedirs(output_dir, exist_ok=True)

        with open(f"{output_dir}/training_losses.json", "w") as f:
            json.dump(trainer.train_losses, f)

        with open(f"{output_dir}/validation_losses.json", "w") as f:
            json.dump(trainer.validation_losses, f)

        return parameters, dataset_size, {}


def client_fn(context: Context):
    """Returns a FlowerClient containing its data partition."""

    initial_model = ModelBuilder().with_model_name(model_name).enable_lora(lora).with_lora_rank(lora_rank).build()
    tokenizer = get_tokenizer(model_name)

    cid = int(context.node_config["partition-id"])

    return LLMClient(cid, initial_model, tokenizer).to_client()

Servidor

def server_fn(context: Context):
    # Objeto callable para criar configuração do cliente
    on_fit_config_fn = FitConfigFactory(train_context_dict)

    # Construção do modelo global
    global_model = ModelBuilder().with_model_name(model_name).enable_lora(lora).with_lora_rank(lora_rank).build()

    # 1 - Entender a inicialização realizada pelo notebook do allan
    # 2 - Entender a manipulação de parâmetros realizada por ela
    initial_ndarrays = get_parameters(global_model, lora)
    initial_parameters = ndarrays_to_parameters(initial_ndarrays)
    evaluate_fn = get_evaluate_fn(eval_context_dict, device)

    # 1 - próximo passo é usar uma injeção de dependências na inversão de contexto.
    strategy = FedAvg(initial_parameters=initial_parameters, fraction_fit=fraction_fit, fraction_evaluate=fraction_eval,
                      on_fit_config_fn=on_fit_config_fn, evaluate_fn=evaluate_fn)

    # Construct ServerConfig
    config = ServerConfig(num_rounds=num_rounds)

    # Wrap everything into a `ServerAppComponents` object
    return ServerAppComponents(strategy=strategy, config=config)

Simulando a Detecção de Anomalias com LLM

client_app = ClientApp(client_fn)
server_app = ServerApp(server_fn=server_fn)

# https://flower.ai/docs/framework/how-to-run-simulations.html
run_simulation(
    server_app=server_app, client_app=client_app, num_supernodes=num_supernodes, backend_config={"client_resources": {"num_cpus": 0.25,"num_gpus": 1}}
)