From 5d03c911328056c133480c60560332eee8bb00dc Mon Sep 17 00:00:00 2001 From: akoumpa Date: Wed, 23 Oct 2024 21:32:28 +0000 Subject: [PATCH] Apply isort and black reformatting Signed-off-by: akoumpa --- examples/llm/peft/hf.py | 22 +++++++++++++--------- examples/llm/sft/hf.py | 25 +++++++++++++++++-------- nemo/collections/llm/peft/lora.py | 13 +++++++++---- 3 files changed, 39 insertions(+), 21 deletions(-) diff --git a/examples/llm/peft/hf.py b/examples/llm/peft/hf.py index 80dc3db513459..86b5199aaa5b7 100644 --- a/examples/llm/peft/hf.py +++ b/examples/llm/peft/hf.py @@ -22,7 +22,8 @@ def mk_hf_dataset(tokenizer): - EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN + EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN + def formatting_prompts_func(examples): alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request. @@ -35,18 +36,22 @@ def formatting_prompts_func(examples): ### Response: {}""" instruction = examples["context"] - input = examples["question"] - output = examples["answers"]['text'] + input = examples["question"] + output = examples["answers"]['text'] if isinstance(output, list): output = output[0] text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN ans = tokenizer(text) tokens = ans['input_ids'] - return {'tokens': tokens, 'labels': tokens[1:] + [tokens[-1]], } + return { + 'tokens': tokens, + 'labels': tokens[1:] + [tokens[-1]], + } from datasets import load_dataset - dataset = load_dataset("rajpurkar/squad", split = "train") - dataset = dataset.map(formatting_prompts_func, batched = False, batch_size = 2) + + dataset = load_dataset("rajpurkar/squad", split="train") + dataset = dataset.map(formatting_prompts_func, batched=False, batch_size=2) return dataset @@ -83,8 +88,7 @@ def formatting_prompts_func(examples): llm.api.finetune( model=llm.HfAutoModelForCausalLM(args.model), data=llm.HfDatasetDataModule( - mk_hf_dataset(tokenizer.tokenizer), - pad_token_id=tokenizer.tokenizer.eos_token_id + mk_hf_dataset(tokenizer.tokenizer), pad_token_id=tokenizer.tokenizer.eos_token_id ), trainer=nl.Trainer( devices=args.devices, @@ -102,4 +106,4 @@ def formatting_prompts_func(examples): optim=fdl.build(llm.adam.pytorch_adam_with_flat_lr(max_lr=1e-5, clip_grad=0.5)), log=None, peft=lora, - ) \ No newline at end of file + ) diff --git a/examples/llm/sft/hf.py b/examples/llm/sft/hf.py index 59a20e5e530fb..e3b51d2a9257c 100644 --- a/examples/llm/sft/hf.py +++ b/examples/llm/sft/hf.py @@ -45,8 +45,10 @@ def _create_dataloader(self, dataset, **kwargs) -> DataLoader: **kwargs1, ) + def mk_hf_dataset(tokenizer): - EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN + EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN + def formatting_prompts_func(examples): alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request. @@ -59,20 +61,25 @@ def formatting_prompts_func(examples): ### Response: {}""" instruction = examples["context"] - input = examples["question"] - output = examples["answers"]['text'] + input = examples["question"] + output = examples["answers"]['text'] if isinstance(output, list): output = output[0] text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN ans = tokenizer(text) tokens = ans['input_ids'] - return {'tokens': tokens, 'labels': tokens[1:] + [tokens[-1]], } + return { + 'tokens': tokens, + 'labels': tokens[1:] + [tokens[-1]], + } from datasets import load_dataset - dataset = load_dataset("rajpurkar/squad", split = "train") - dataset = dataset.map(formatting_prompts_func, batched = False, batch_size = 2) + + dataset = load_dataset("rajpurkar/squad", split="train") + dataset = dataset.map(formatting_prompts_func, batched=False, batch_size=2) return dataset + def squad(tokenizer) -> pl.LightningDataModule: return SquadDataModuleWithMbs( tokenizer=tokenizer, @@ -83,13 +90,16 @@ def squad(tokenizer) -> pl.LightningDataModule: sanity_check_dist_workers=False, ) + class HfAutoModelPeft(llm.HfAutoModel): def configure_model(self): super().configure_model() self.model.eval() from lora import apply_lora_to_model + apply_lora_to_model(self.model) + if __name__ == '__main__': import argparse @@ -119,8 +129,7 @@ def configure_model(self): llm.api.finetune( model=HfAutoModelPeft(args.model), data=llm.HfDatasetDataModule( - mk_hf_dataset(tokenizer.tokenizer), - pad_token_id=tokenizer.tokenizer.eos_token_id + mk_hf_dataset(tokenizer.tokenizer), pad_token_id=tokenizer.tokenizer.eos_token_id ), trainer=nl.Trainer( devices=args.devices, diff --git a/nemo/collections/llm/peft/lora.py b/nemo/collections/llm/peft/lora.py index e2ca4a8cdb818..855014666b2da 100644 --- a/nemo/collections/llm/peft/lora.py +++ b/nemo/collections/llm/peft/lora.py @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -import re -import torch import math +import re from dataclasses import dataclass, field from typing import List, Literal +import torch from megatron.core import parallel_state from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear from torch import nn @@ -72,7 +72,9 @@ def forward(self, x): class LinearAdapter(nn.Module): - def __init__(self, orig_linear, dim=8, alpha=32, dropout=0.1, dropout_position='post', lora_A_init_method='xavier'): + def __init__( + self, orig_linear, dim=8, alpha=32, dropout=0.1, dropout_position='post', lora_A_init_method='xavier' + ): super(LinearAdapter, self).__init__() assert isinstance(orig_linear, nn.Linear) @@ -111,6 +113,7 @@ def forward(self, x): lora_res = self.dropout(lora_res) return res + lora_res + @dataclass class LoRA(PEFT): """ @@ -212,7 +215,9 @@ def wildcard_match(pattern, key): in_features = m.input_size out_features = m.output_size elif isinstance(m, nn.Linear): - return LinearAdapter(m, dim=self.dim, alpha=self.alpha, dropout=self.dropout, lora_A_init_method=self.lora_A_init_method) + return LinearAdapter( + m, dim=self.dim, alpha=self.alpha, dropout=self.dropout, lora_A_init_method=self.lora_A_init_method + ) else: raise NotImplementedError(f"Layer type is unrecognized for LoRA: {type(m)}")