LLM Training and Safety Alignment: A Practical Guide for Developers

Start with the compute bottleneck: choosing a parallelism strategy

Training a model with hundreds of billions of parameters means a single GPU can neither hold the model nor process enough data. The goal of distributed training is to make multiple GPUs work simultaneously while keeping communication overhead acceptable. The core trade-off is between compute efficiency and communication volume.

The three common parallelism strategies are as follows.

Strategy	What is split	Communication characteristics	When to use
Data parallelism	Each GPU holds the full model and splits the training batch	Synchronizes gradients at each step; communication scales with parameter count	The model fits in single-GPU memory
Tensor parallelism	Splits weight matrices within a single layer	Frequent intra-layer communication; requires high-bandwidth interconnect	A single layer is too large; often used within the same node
Pipeline parallelism	Splits the model by layers across devices	Passes only activations, but has pipeline bubbles	Many layers; needs cross-node scaling

Real projects usually combine these strategies. For example, tensor parallelism within a node, pipeline parallelism across nodes, and data parallelism on the outside. A practical decision order is: first check whether a single GPU can hold model parameters, gradients, and optimizer states. If it can, prefer data parallelism. If not, consider sharding approaches such as ZeRO or FSDP, which distribute parameters, gradients, and optimizer states. If that is still insufficient, introduce tensor parallelism or pipeline parallelism.

Adding compute does not always yield a proportional training speedup. Kaplan et al. (2020) and Hoffmann et al. (2022) show empirical power-law relationships among model parameter count, training data volume, and compute budget. When data is insufficient or the budget allocation is unbalanced, the returns from adding more compute drop noticeably. These conclusions come from specific architectures and data configurations; validate them with small-scale experiments before applying them to a new model.

Minimal runnable DDP example

The following code uses PyTorch’s DistributedDataParallel. The runtime requires PyTorch with CUDA support, NVIDIA GPUs, and the NCCL communication backend. The model is a placeholder implementation, and the data is randomly generated; it is only used to verify whether the distributed workflow runs.

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
class TinyLM(torch.nn.Module):
    """Placeholder model. No causal mask is applied; this is only for validating the training workflow. Replace it with your own Transformer implementation in real projects."""
    def __init__(self, vocab_size=32000, hidden=1024, layers=4):
        super().__init__()
        self.embed = torch.nn.Embedding(vocab_size, hidden)
        block = torch.nn.TransformerEncoderLayer(
            d_model=hidden, nhead=16, batch_first=True
        )
        self.encoder = torch.nn.TransformerEncoder(block, num_layers=layers)
        self.head = torch.nn.Linear(hidden, vocab_size)
    def forward(self, input_ids):
        x = self.encoder(self.embed(input_ids))
        return self.head(x)
def main():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    model = TinyLM().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
    loss_fn = torch.nn.CrossEntropyLoss()
    # Random data is only for validating the workflow; replace it with tokenized corpora for real training
    tokens = torch.randint(0, 32000, (2048, 512))
    dataset = TensorDataset(tokens)
    sampler = DistributedSampler(dataset, shuffle=True)
    loader = DataLoader(dataset, batch_size=4, sampler=sampler)
    model.train()
    for epoch in range(1):
        sampler.set_epoch(epoch)  # Use a different random order for each epoch
        for (batch,) in loader:
            batch = batch.cuda(local_rank)
            logits = model(batch[:, :-1])
            loss = loss_fn(
                logits.reshape(-1, logits.size(-1)),
                batch[:, 1:].reshape(-1),
            )
            optimizer.zero_grad()
            loss.backward()  # DDP performs the gradient all-reduce here
            optimizer.step()
        if dist.get_rank() == 0:
            print(f"epoch {epoch} last loss {loss.item():.4f}")
    dist.destroy_process_group()
if __name__ == "__main__":
    main()

The launch command for a single node with 8 GPUs is:

torchrun --nproc_per_node=8 train.py

torchrun automatically sets environment variables such as LOCAL_RANK, RANK, and WORLD_SIZE. Multi-node training also requires —nnodes, —node_rank, and —rdzv_endpoint. Real training must use DistributedSampler to split data. If every GPU reads the same samples, you are effectively duplicating computation and wasting compute.

Common troubleshooting points

• Out of memory (OOM): reduce the batch size first, then enable activation checkpointing, and finally consider sharding approaches.
• Slow training: use a profiler to check the share of time spent on communication. If all-reduce takes up too high a share, check network bandwidth, the NCCL version, and gradient accumulation steps.
• Multi-GPU results differ from single-GPU results: check random seeds, data splitting, and dropout settings on each GPU.

Safety alignment: workflow and limitations

As model capabilities increase, output risks also rise. The goal of alignment is to strike a suitable balance between helpfulness and safety, not to make the model reject as many requests as possible. Excessive refusal harms usability, while unchecked output creates compliance and safety issues.

A practical alignment workflow usually includes the following steps:

1. Define behavior specifications: clarify which requests must be refused, which require additional warnings, and which can be answered directly.
2. Supervised fine-tuning (SFT): use human-labeled demonstration data to train the model to follow instruction formats and boundaries.
3. Preference optimization: use methods such as RLHF and DPO to help the model learn response styles that humans prefer.
4. Red teaming: use human or automated tools to construct jailbreak prompts, leading questions, and edge cases.
5. Evaluation and staged rollout: track refusal rate, false refusal rate, and jailbreak success rate; launch to a small traffic share first, then expand gradually.
6. Audit and rollback: retain sanitized interaction logs and be able to quickly roll back to the previous version when issues arise.

For example, in a medical consultation scenario, the system should identify high-risk questions such as medication dosage and symptom assessment, provide general health information, and recommend that users consult a licensed physician. It should not replace diagnosis.