Integrate custom models with TorchAcc

TorchAcc optimization methods

TorchAcc offers two main optimization methods. Understanding them will help you choose the best approach to improve your model's training speed and efficiency.

• Compilation optimization

  TorchAcc converts a PyTorch dynamic graph into a static graph, which it then optimizes and compiles. A JIT compiler transforms the computation graph into more efficient code. This process avoids the performance overhead common in dynamic graph execution and improves training speed and efficiency.

• Custom optimization

  Global compilation optimization cannot accelerate distributed training for models with dynamic shapes, custom operators, or dynamic control flow. For these models, TorchAcc provides custom optimizations such as IO optimization, computation (kernel) optimization, and memory optimization.

TorchAcc compilation optimization

Integrate with distributed training

Follow these steps to integrate the TorchAcc compiler for distributed training:

1. Set a fixed random seed.

   Setting a fixed random seed ensures that weights are initialized consistently across all workers. This is an alternative to broadcasting weights.

   torch.manual_seed(SEED_NUMBER)
   # Replace with:
   xm.set_rng_state(SEED_NUMBER)
   

2. After you obtain the xla_device, call set_replication, wrap the dataloader, and set the model's device placement.

   device = xm.xla_device()
   xm.set_replication(device, [device])
   # Wrapper dataloader
   data_loader_train = pl.MpDeviceLoader(data_loader_train, device)
   data_loader_val = pl.MpDeviceLoader(data_loader_val, device)
   # Dispatch device to model
   model.to(device)

3. Initialize the distributed environment.

   Set the backend parameter of dist.init_process_group to 'xla':

   dist.init_process_group(backend='xla', init_method='env://')

4. All-reduce gradients.

   After the backward pass on the loss, all-reduce the gradients:

   gradients=xm._fetch_gradients(optimizer)
   xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())

   Important

   When you train with automatic mixed precision (AMP), make sure to call xm.all_reduce before you call scaler.unscale_. This ensures that overflow detection runs on the reduced gradients.

5. Launch the job using xlarun.

   xlarun --nproc_per_node=8 YOUR_MODEL.py

   Note

   For multi-node jobs, the usage is the same as torchrun.

Integrate with mixed precision

You can use mixed precision training to accelerate your model. The following steps add automatic mixed precision (AMP) to the compilation optimization setup from the previous section.

1. Implement AMP by using the native PyTorch functionality.

   Using mixed precision in TorchAcc is nearly identical to its native PyTorch implementation. First, refer to the following documents to implement native PyTorch AMP.

   • AUTOMATIC MIXED PRECISION PACKAGE - TORCH.AMP

   • CUDA AUTOMATIC MIXED PRECISION EXAMPLES

   • AUTOMATIC MIXED PRECISION

2. Replace GradScaler.

   Replace torch.cuda.amp.GradScaler with torchacc.torch_xla.amp.GradScaler:

   from torchacc.torch_xla.amp import GradScaler

3. Replace the optimizer.

   Native PyTorch optimizers may result in slightly lower performance. To further improve training speed, you can replace an optimizer from torch.optim with a syncfree optimizer.

   from torchacc.torch_xla.amp import syncfree
   adam_optimizer = syncfree.Adam()
   adamw_optimizer = syncfree.AdamW()
   sgd_optimizer = syncfree.SGD()

   Currently, the syncfree optimizer module provides only these three optimizers. For all other types, continue to use the native PyTorch optimizers.

Example

The following code shows an example with a Bert-base model:

import argparse
import os
import time
import torch
import torch.distributed as dist
from datasets import load_from_disk
from datetime import datetime as dt
from time import gmtime, strftime
from transformers import AutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding
# Pytorch1.12 default set False.
torch.backends.cuda.matmul.allow_tf32=True
parser = argparse.ArgumentParser()
parser.add_argument("--amp-level", choices=["O1"], default="O1", help="amp level.")
parser.add_argument("--profile", action="store_true")
parser.add_argument("--profile_folder", type=str, default="./profile_folder")
parser.add_argument("--dataset_path", type=str, default="./sst_data/train")
parser.add_argument("--num_epochs", type=int, default=1)
parser.add_argument("--batch_size", type=int, default=128)
parser.add_argument("--max_seq_length", type=int, default=512)
parser.add_argument("--break_step_for_profiling", type=int, default=20)
parser.add_argument("--model_name", type=str, default="bert-base-cased")
parser.add_argument("--local_rank", type=int, default="-1")
parser.add_argument("--log-interval", type=int, default="10")
parser.add_argument('--max-steps', type=int, default=200, help='total training epochs.')
args = parser.parse_args()
print("Job running args: ", args)
args.local_rank = os.getenv("LOCAL_RANK", 0)
+def enable_torchacc_compiler():
+    return os.getenv('TORCHACC_COMPILER_OPT') is not None
def print_rank_0(message):
    """If distributed is initialized, print only on rank 0."""
    if torch.distributed.get_rank() == 0:
        print(message, flush=True)
def print_test_update(epoch, step, batch_size, loss, time_elapsed, samples_per_step, peak_mem):
  # Getting the current date and time
  dt = strftime("%a, %d %b %Y %H:%M:%S", gmtime())
  print_rank_0(train_format_string.format(dt, epoch, step, batch_size, loss, time_elapsed, samples_per_step, peak_mem))
def log_metrics(epoch, step, batch_size, loss, batch_time, samples_per_step, peak_mem):
    batch_time = f"{batch_time:.3f}"
    samples_per_step = f"{samples_per_step:.3f}"
    peak_mem = f"{peak_mem:.3f}"
+    if enable_torchacc_compiler():
+        import torchacc.torch_xla.core.xla_model as xm
+        xm.add_step_closure(
+            print_test_update, args=(epoch, step, batch_size, loss, batch_time, samples_per_step, peak_mem), run_async=True)
+    else:
        print_test_update(epoch, step, batch_size, loss, batch_time, samples_per_step, peak_mem)
+if enable_torchacc_compiler():
+  import torchacc.torch_xla.core.xla_model as xm
+  import torchacc.torch_xla.distributed.parallel_loader as pl
+  import torchacc.torch_xla.distributed.xla_backend
+  from torchacc.torch_xla.amp import autocast, GradScaler, syncfree
+  xm.set_rng_state(101)
+  dist.init_process_group(backend="xla", init_method="env://")
+else:
  from torch.cuda.amp import autocast, GradScaler
  dist.init_process_group(backend="nccl", init_method="env://")
dist.barrier()
args.world_size = dist.get_world_size()
args.rank = dist.get_rank()
print("world size:", args.world_size, " rank:", args.rank, " local rank:", args.local_rank)
def get_autocast_and_scaler():
+  if enable_torchacc_compiler():
+    return autocast, GradScaler()
  return autocast, GradScaler()
def loop_with_amp(model, inputs, optimizer, autocast, scaler):
  with autocast():
    outputs = model(**inputs)
    loss = outputs["loss"]
  scaler.scale(loss).backward()
+  if enable_torchacc_compiler():
+    gradients = xm._fetch_gradients(optimizer)
+    xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())
  scaler.step(optimizer)
  scaler.update()
  return loss, optimizer
def loop_without_amp(model, inputs, optimizer):
  outputs = model(**inputs)
  loss = outputs["loss"]
  loss.backward()
+  if enable_torchacc_compiler():
+    xm.optimizer_step(optimizer)
+  else:
    optimizer.step()
  return loss, optimizer
def full_train_epoch(epoch, model, train_device_loader, device, optimizer, autocast, scaler, profiler=None):
  model.train()
  iteration_time = time.time()
  num_steps = int(len(train_device_loader.dataset) / args.batch_size)
  for step, inputs in enumerate(train_device_loader):
    if step > args.max_steps:
      break
+   if not enable_torchacc_compiler():
      inputs.to(device)
    optimizer.zero_grad()
    if args.amp_level == "O1":
      loss, optimizer = loop_with_amp(model, inputs, optimizer, autocast, scaler)
    else:
      loss, optimizer = loop_without_amp(model, inputs, optimizer)
    if args.profile and profiler:
      profiler.step()
    if step % args.log_interval == 0:
      time_elapsed = (time.time() - iteration_time) / args.log_interval
      iteration_time = time.time()
      samples_per_step = float(args.batch_size / time_elapsed) * args.world_size
      peak_mem = torch.cuda.memory_allocated()/1024.0/1024.0/1024.0
      log_metrics(epoch, step, args.batch_size, loss, time_elapsed, samples_per_step, peak_mem)
def train_bert():
  model = AutoModelForSequenceClassification.from_pretrained(args.model_name, cache_dir="./model")
  tokenizer = AutoTokenizer.from_pretrained(args.model_name)
  tokenizer.model_max_length = args.max_seq_length
  training_dataset = load_from_disk(args.dataset_path)
  collator = DataCollatorWithPadding(tokenizer)
  training_dataset = training_dataset.remove_columns(['text'])
  train_device_loader = torch.utils.data.DataLoader(
      training_dataset, batch_size=args.batch_size, collate_fn=collator, shuffle=True, num_workers=4)
+  if enable_torchacc_compiler():
+    device = xm.xla_device()
+    xm.set_replication(device, [device])
+    train_device_loader = pl.MpDeviceLoader(train_device_loader, device)
+    model = model.to(device)
+  else:
    device = torch.device(f"cuda:{args.local_rank}")
    torch.cuda.set_device(device)
    model = model.cuda()
    model = torch.nn.parallel.DistributedDataParallel(model)
+  if enable_torchacc_compiler() and args.amp_level == "O1":
+    optimizer = syncfree.Adam(model.parameters(), lr=1e-3)
+  else:
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
  autocast, scaler = None, None
  if args.amp_level == "O1":
    autocast, scaler = get_autocast_and_scaler()
  if args.profile:
    with torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=2, warmup=2, active=20),
        on_trace_ready=torch.profiler.tensorboard_trace_handler(args.profile_folder)) as prof:
      for epoch in range(args.num_epochs):
        full_train_epoch(epoch, model, train_device_loader, device, optimizer, autocast, scaler, profiler=prof)
  else:
    for epoch in range(args.num_epochs):
      full_train_epoch(epoch, model, train_device_loader, device, optimizer, autocast, scaler)
if __name__ == "__main__":
  train_bert()

TorchAcc custom optimizations

IO optimization

Data prefetcher

The Prefetcher prefetches training data and supports data preprocessing through the preprocess_fn parameter.

+ from torchacc.runtime.io.prefetcher import Prefetcher
data_loader = build_data_loader()
model = build_model()
optimizer = build_optimizer()
# define preprocess function
preprocess_fn = None
+ prefetcher = Prefetcher(data_loader, preprocess_fn)
for iter, samples in enumerate(prefetcher):
    loss = model(samples)
    loss.backward()
    # Prefetch to CPU first. Call after backward and before update.
    # At this point we are waiting for kernels launched by cuda graph
    #  to finish, so CPU is idle. Take advantage of this by loading next
    #  input batch before calling step.
+    prefetcher.prefetch_CPU()
    optimizer.step()
    # Prefetch to GPU. Call after optimizer step.
+	prefetcher.prefetch_GPU()

Pack dataset

This method packs variable-length samples from language datasets into fixed-shape batches. This reduces the proportion of zero-padding and minimizes the dynamic nature of batch data, improving training efficiency.

Pin memory

When you define the dataloader, add the pin_memory parameter and increase num_workers as needed.

train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=self.batch_size_per_device,
            num_workers=8,
            pin_memory=True,
            drop_last=True,
            shuffle=False,
            collate_fn=partial(collate_fn))

Computation optimization

Kernel fusion optimizations

The following fusion optimizations are available:

• FusedLayerNorm

  # An equivalent replacement kernel for LayerNorm.
  from torchacc.runtime import hooks
  # add before import torch
  hooks.enable_fused_layer_norm()

• FusedAdam

  # An equivalent replacement kernel for Adam/AdamW.
  from torchacc.runtime import hooks
  # add before import torch
  hooks.enable_fused_adam()

• QuickGelu

  # Replaces nn.GELU with QuickGelu.
  from torchacc.runtime.nn.quick_gelu import QuickGelu

• fused_bias_dropout_add

  # from torchacc.runtime.nn import dropout_add_fused_train, 
  # Fuses operations such as Dropout and element-wise bias add.
  if self.training:
      # train mode
      with torch.enable_grad():
          x = dropout_add_fused_train(x, to_add, drop_rate)
  else:
      # inference mode
      x = dropout_add_fused(x, to_add, drop_rate)

• WindowProcess

  # The WindowProcess optimization kernel fuses operations related to window shifting and partitioning in SwinTransformer, including:
  # - window cyclic shift and window partition
  # - window merge and reverse cyclic shift
  from torchacc.runtime.nn.window_process import WindowProcess
  if not fused:
      shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
      x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
  else:
      x_windows = WindowProcess.apply(x, B, H, W, C, -self.shift_size, self.window_size)
  from torchacc.runtime.nn.window_process import WindowProcessReverse
  if not fused:
      shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
      x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
  else:
      x = WindowProcessReverse.apply(attn_windows, B, H, W, C, self.shift_size, self.window_size)

• FusedSwinFmha

  # Fuses the qk_result + relative_position_bias + mask + softmax part of the Multi-Head Attention (MHA) in SwinTransformer.
  from torchacc.runtime.nn.fmha import FusedSwinFmha
  FusedSwinFmha.apply(attn, relative_pos_bias, attn_mask, batch_size, window_num,
                num_head, window_len)

• nms/nms_normal/soft_nms/batched_soft_nms

  # Fused CUDA kernel implementations for four types of operators: nms, nms_normal, soft_nms, and batched_soft_nms.
  from torchacc.runtime.nn.nms import nms, nms_normal
  from torchacc.runtime.nn.nms import soft_nms, batched_soft_nms
        

Parallelized kernel optimizations

DCN/DCNv2:

# Parallelizes computation for the dcn_v2_cuda backward pass.
from torchacc.runtime.nn.dcn_v2 import DCN, DCNv2
self.conv = DCN(chi, cho, kernel_size, stride, padding, dilation, deformable_groups)

Multi-stream kernel optimizations

This optimization uses multiple streams to compute a set of inputs for a function concurrently.

from torchacc.runtime.utils.misc import multi_apply_multi_stream
from mmdet.core import multi_apply
def test_func(t1, t2, t3):
  t1 = t1 * 2.0
  t2 = t2 + 2.0
  t3 = t3 - 2.0
  return (t1, t2, t3)
cuda = torch.device('cuda')
t1 = torch.empty((100, 1000), device=cuda).normal_(0.0, 1.0)
t2 = torch.empty((100, 1000), device=cuda).normal_(0.0, 2.0)
t3 = torch.empty((100, 1000), device=cuda).normal_(0.0, 3.0)
if enable_torchacc:
    result = multi_apply_multi_stream(test_func, 2, t1, t2, t3)
else:
    result = multi_apply(test_func, t1, t2, t3)

Memory optimization

Gradient checkpointing

import torchacc
model = torchacc.auto_checkpoint(model)