Skip to content

Instantly share code, notes, and snippets.

@woshiyyya
Last active August 16, 2023 01:22
Show Gist options
  • Save woshiyyya/6191aebbe0c2edce9b7ba382676ab421 to your computer and use it in GitHub Desktop.
Save woshiyyya/6191aebbe0c2edce9b7ba382676ab421 to your computer and use it in GitHub Desktop.
# Minimal Example adapted from https://huggingface.co/docs/transformers/training
import deepspeed
import evaluate
import torch
from datasets import load_dataset
from deepspeed.accelerator import get_accelerator
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
AutoModelForSequenceClassification,
AutoTokenizer,
set_seed,
)
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
def get_datasets(tokenizer):
"""Creates and preprocess the GLUE MRPC dataset."""
datasets = load_dataset("glue", "mrpc")
def tokenize_function(examples):
outputs = tokenizer(
examples["sentence1"],
examples["sentence2"],
truncation=True,
max_length=None,
)
return outputs
tokenized_datasets = datasets.map(
tokenize_function,
batched=True,
remove_columns=["idx", "sentence1", "sentence2"],
)
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
return tokenized_datasets
def train_func(config):
# Set global random seed
set_seed(config["seed"])
# Load datasets and metrics
metric = evaluate.load("glue", "mrpc")
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
tokenized_datasets = get_datasets(tokenizer)
# Create evaluation dataloader
# The training dataloader will be created by `deepspeed.initialize` later
def collate_fn(examples):
return tokenizer.pad(
examples,
padding="longest",
return_tensors="pt",
)
eval_dataloader = DataLoader(
tokenized_datasets["validation"],
batch_size=config["eval_batch_size"],
collate_fn=collate_fn,
shuffle=False,
)
# Instantiate the model
model = AutoModelForSequenceClassification.from_pretrained(
"bert-base-cased", return_dict=True
)
# Initialize DeepSpeed Engine
model, optimizer, train_dataloader, lr_scheduler = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
training_data=tokenized_datasets["train"],
collate_fn=collate_fn,
config=deepspeed_config,
)
# Define the GPU device for the current worker
device = get_accelerator().device_name(model.local_rank)
# Now we train the model
for epoch in range(config["num_epochs"]):
model.train()
for batch in tqdm(
train_dataloader,
desc=f"Training epoch {epoch}",
disable=(model.global_rank != 0),
):
batch = {k: v.to(device) for k, v in batch.items()}
outputs = model(**batch)
loss = outputs.loss
model.backward(loss)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
model.eval()
for batch in tqdm(
eval_dataloader,
desc=f"Evaluation epoch {epoch}",
disable=(model.global_rank != 0),
):
batch = {k: v.to(device) for k, v in batch.items()}
with torch.no_grad():
outputs = model(**batch)
predictions = outputs.logits.argmax(dim=-1)
metric.add_batch(
predictions=predictions,
references=batch["labels"],
)
eval_metric = metric.compute()
if model.global_rank == 0:
print(f"epoch {epoch}:", eval_metric)
if __name__ == "__main__":
deepspeed_config = {
"optimizer": {
"type": "AdamW",
"params": {
"lr": 2e-5,
},
},
"scheduler": {"type": "WarmupLR", "params": {"warmup_num_steps": 100}},
"fp16": {"enabled": True},
"bf16": {"enabled": False}, # Turn this on if using AMPERE GPUs.
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "none",
},
"offload_param": {
"device": "none",
},
},
"gradient_accumulation_steps": 1,
"gradient_clipping": True,
"steps_per_print": 10,
"train_micro_batch_size_per_gpu": 16,
"wall_clock_breakdown": False,
}
training_config = {
"seed": 42,
"num_epochs": 3,
"eval_batch_size": 32,
"deepspeed_config": deepspeed_config,
}
trainer = TorchTrainer(
train_func,
train_loop_config=training_config,
scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
)
trainer.fit()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment