[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[project]
name = "flowertune-llm"
version = "1.0.0"
description = "FlowerTune LLM: Federated LLM Fine-tuning with Flower"
license = "Apache-2.0"
dependencies = [
    "flwr[simulation]==1.12.0",
    "flwr-datasets>=0.3.0",
    "trl==0.8.1",
    "bitsandbytes==0.43.0",
    "scipy==1.13.0",
    "peft==0.6.2",
    "fschat[model_worker,webui]==0.2.35",
    "transformers==4.41.1",
    "sentencepiece==0.2.0",
    "omegaconf==2.3.0",
    "hf_transfer==0.1.8",
]


[tool.hatch.build.targets.wheel]
packages = ["."]

[tool.flwr.app]
publisher = "flwrlabs"

[tool.flwr.app.components]
serverapp = "fedllm.server_app:app"
clientapp = "fedllm.client_app:app"

[tool.flwr.app.config]
num-server-rounds = 2
num-supernodes = 10

# Define dataset
dataset.type = 'homo' # type = ['homo','hete']
dataset.name = "vicgalle/alpaca-gpt4"

# Define model settings
model.name = "Qwen/Qwen2.5-1.5B-Instruct" 
model.quantization = 4
model.gradient-checkpointing = true
model.flash-attention = false

### Use MATES ###
mates.state = true
mates.holdout-ratio = 0.001
mates.reference-ratio = 0.0005
mates.holdout-batch-size = 4
mates.reference-batch-size = 2
mates.update-data-influence-model-step = 20
mates.selection-fraction = 0.4
### END ###

### Use SkipBERT ###

# Model setting
skipbert.student-model = "bert-base-uncased"
skipbert.output-mode = "regression"
skipbert.num-layers-student = 12
skipbert.num-full-hidden-layers-student = 6
skipbert.num-masked-layers-teacher = 0
skipbert.num-masked-last-layers-teacher = 0

# Training hyperparameters
skipbert.train-batch-size = 4
skipbert.gradient-accumulation-steps = 1
skipbert.eval-batch-size = 4
skipbert.eval-accumulation-steps = 1
skipbert.learning-rate = 2.0e-5
skipbert.num-train-epochs = 10
skipbert.eval-step = 10
skipbert.max-seq-length = 256
skipbert.weight-decay = 1.0e-4
skipbert.warmup-steps = 100 # 500
skipbert.do-train = true
skipbert.do-eval = true
skipbert.do-predict = false
skipbert.max-steps = -1
skipbert.evaluation-strategy = "epoch"
skipbert.save-strategy = "epoch"
skipbert.lr-scheduler-type = "cosine" # or 'warmup-linear'
skipbert.logging-dir = './skipbert-logs'
skipbert.output-dir = "./skipbert-results"
skipbert.report-to = 'wandb'

# Knowledge distillation parameters
skipbert.beta = 0.01
skipbert.T = 1.0
skipbert.alpha = 1.0
skipbert.reduce-T = 1.0
skipbert.epochs-no-cls = 5

# Training schedule and features
skipbert.freeze-lower-layers = true

# Feature usage flags
skipbert.use-logits = true
skipbert.use-att = true
skipbert.use-rep = true
skipbert.use-embedding = false

# Training modes
skipbert.do-fit = false
skipbert.fp16 = false
skipbert.no-pretrain = false
skipbert.use-init-weight = false
skipbert.share-param = true
skipbert.do-lower-case = true
skipbert.no-cuda = false

# N-gram settings
skipbert.n-gram-left = 1
skipbert.n-gram-right = 1

# Layer mappings
skipbert.att-layer-maps = "1, 3, 5, 7, 9, 11"
skipbert.hid-layer-maps = "6, 7, 8, 9, 10, 11, 12"

### END ###




# Define LoRA settings
model.lora.lora-r = 8 
model.lora.lora-alpha = 16
model.lora.lora-dropout = 0.05
model.lora.lora-target-modules = "q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj"

# Define training settings
train.save-every-round = 5 
train.learning-rate-max = 5e-5
train.learning-rate-min = 1e-6
train.seq-length = 256
train.prompt-template-name = "alpaca"
train.train-on-inputs = true
train.verbose = false

# Define training agruments for HF Trainer
train.training-arguments.output-dir = "" 
train.training-arguments.learning-rate = 3e-4
train.training-arguments.per-device-train-batch-size = 4
train.training-arguments.gradient-accumulation-steps = 1
train.training-arguments.per-device-eval-batch-size = 2
train.training-arguments.eval-accumulation-steps = 1
train.training-arguments.logging-steps = 10
train.training-arguments.num-train-epochs = 1
train.training-arguments.max-steps = 10
train.training-arguments.save-steps = 1000
train.training-arguments.save-total-limit = 10
train.training-arguments.gradient-checkpointing = true
train.training-arguments.lr-scheduler-type = "cosine"
train.training-arguments.warmup-steps = 0
train.training-arguments.do-train = true
train.training-arguments.do-eval = true
train.training-arguments.dataloader-drop-last = false
train.training-arguments.eval-strategy = "epoch"
train.training-arguments.save-strategy = "epoch"
train.training-arguments.ddp-find-unused-parameters = false
train.training-arguments.group-by-length = true
train.training-arguments.load-best-model-at-end = true
train.training-arguments.report-to = "wandb"

# Define local training settings 
train.strategy.fraction-fit = 0.2 
train.strategy.fraction-evaluate = 0.0

[tool.flwr.federations]
default = "local-simulation"

[tool.flwr.federations.local-simulation]
options.num-supernodes = 10
options.backend.client-resources.num-cpus = 8
options.backend.client-resources.num-gpus = 1.0