Hugging Face Accelerate, Weights & Biases로 초강력 업그레이드

pip install accelerate
OR
conda install -c conda-forge accelerate

pip install git+https://github.com/huggingface/accelerate

accelerate config

accelerate env

- `Accelerate` version: 0.11.0.dev0
- Platform: Linux-5.10.0-15-cloud-amd64-x86_64-with-debian-11.3
- Python version: 3.7.12
- Numpy version: 1.19.5
- PyTorch version (GPU?): 1.12.0+cu102 (True)
- `Accelerate` default config:
        - compute_environment: LOCAL_MACHINE
        - distributed_type: MULTI_GPU
        - mixed_precision: fp16
        - use_cpu: False
        - num_processes: 2
        - machine_rank: 0
        - num_machines: 1
        - main_process_ip: None
        - main_process_port: None
        - main_training_function: main
        - deepspeed_config: {}
        - fsdp_config: {}

 Which type of machine are you using?

compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MPS
downcast_bf16: 'no'
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 1
use_cpu: false

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    optimizer.step()
    scheduler.step()

from accelerate import Accelerator
accelerator = Accelerator()

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
     model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()

from accelerate import Accelerator

accelerator = Accelerator()

from accelerate import Accelerator

accelerator = Accelerator(gradient_accumulation_steps=2)

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
     model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    # accumulate context manager
    with accelerate.accumulate(model):
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        scheduler.step()

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()

        # Unscales the gradients of optimizer's assigned params in-place
        scaler.unscale_(optimizer)

        # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        # optimizer's gradients are already unscaled, so scaler.step does not unscale them,
        # although it still skips optimizer.step() if the gradients contain infs or NaNs.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        scaler.update()

from accelerate import Accelerator

max_grad_norm = 1.0
accelerator = Accelerator(mixed_precision="fp16", gradient_accumulation_steps=2)

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
     model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    # accumulate context manager (for gradient accumulation)
    with accelerate.accumulate(model):
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
	# gradient clipping
	if accelerator.sync_gradients:
            accelerator.clip_grad_norm_(model.parameters(), max_grad_norm)
        optimizer.step()
        scheduler.step()

for inputs, targets in validation_dataloader:
    predictions = model(inputs)
    # Gather all predictions and targets
    all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))
    metrics = calculate_metrics(all_predictions, all_targets)

if accelerator.is_local_main_process:
    do_thing_once_per_server()

@accelerator.on_local_main_process
def do_my_thing():
    "Something done once per server"
    do_thing_once_per_server()

if accelerator.is_main_process:
    do_thing_once()

@accelerator.on_main_process
def do_my_thing():
    "Something done once per server"
    do_thing_once()

@accelerator.on_local_process(local_process_idx=0)
def do_my_thing():
    "Something done on process index 0 on each server"
    do_thing_on_index_zero_on_each_server()

@accelerator.on_process(process_index=0)
def do_my_thing():
    "Something done on process index 0"
    do_thing_on_index_zero()

accelerator.print("My thing I want to print!") 

accelerator.wait_for_everyone()

model = MyModel()
model = accelerator.prepare(model)
accelerator.wait_for_everyone()
# Unwrap
model = accelerator.unwrap_model(model)
state_dict = model.state_dict()
# Use accelerator.save()
accelerator.save(state_dict, "my_state.pkl")

from accelerate import Accelerator
import torch

accelerator = Accelerator()

my_scheduler = torch.optim.lr_scheduler.StepLR(my_optimizer, step_size=1, gamma=0.99)
my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)

# Register the LR scheduler
accelerate.register_for_checkpointing(my_scheduler)

# Save the starting state
accelerate.save_state("my/save/path")

# Perform training
# training loop here ...

# Restore previous state
accelerate.load_state("my/save/path")

from accelerate.logging import get_logger

logger = get_logger(__name__)

# logs on all processes
logger.info("My log", main_process_only=False)
# logs only on main process
logger.debug("My log", main_process_only=True)

from accelerate import Accelerator
accelerator = Accelerator(log_with="wandb")

from accelerate import Accelerator

accelerator = Accelerator()
hps = {"num_epochs": 5, "learning_rate": 1e-4, "batch_size": 16}
accelerator.init_trackers(
    "my_project",
    config=hps,
    init_kwargs={
        "wandb": {
            "notes": "testing accelerate pipeline",
            "tags": ["tag_a", "tag_b"],
            "entity": "gladiator",
        }
    },
)

accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)

accelerator.end_training()

from accelerate import Accelerator

def main():
    accelerator = Accelerator()

    model, optimizer, training_dataloader, scheduler = accelerator.prepare(
        model, optimizer, training_dataloader, scheduler
    )

    for batch in training_dataloader:
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        scheduler.step()


if __name__ == "__main__":
    main()

def main():
   function_which_does_data_processing()
   function_which_does_training()
   function_which_does_evaluation()

if __name__ == "__main__":
    main()

accelerate launch {script_name.py} --arg1 --arg2 ...

CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...

from accelerate import notebook_launcher
# the arguments that the main function takes
args = ("fp16", 42, 64)
notebook_launcher(main, args, num_processes=2)