AIInfra

https://infrasys-ai.github.io/aiinfra-docs/
Admin
AIInfra is an open-source educational curriculum and full-stack system design project covering AI...
Tokens:2,440,633
Snippets:8,553
Trust Score:9.7
Update:1 week ago
Show doc for...
Tokens:
Raw
### Model and optimizer setup with timers and metrics

Source: https://infrasys-ai.github.io/aiinfra-docs/_sources/04Train02ParallelAdv/08PPInterleaved.md

Starts a timer and records a metric before calling setup_model_and_optimizer to build the model, optimizer, and learning rate scheduler.

```python
timers('model-and-optimizer-setup', log_level=0).start(barrier=True)
app_metrics['app_build_optimizer_start_time'] = one_logger_utils.get_timestamp_in_ms()
model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
    model_provider, model_type, checkpointing_context=checkpointing_context
)
```

--------------------------------

### Install required libraries for Qwen3-4B fine-tuning

Source: https://infrasys-ai.github.io/aiinfra-docs/04Train05FineTune/Code03PEFT.html

Installs transformers, peft, datasets, torch, accelerate, and unsloth via pip. This is the first step in the fine-tuning setup.

```bash
# 安装所需库
!pip install transformers peft datasets torch accelerate unsloth
```

--------------------------------

### _MultiProcessingDataLoaderIter.__init__: worker process setup

Source: https://infrasys-ai.github.io/aiinfra-docs/04Train01ParallelBegin/04.data_parallel_implement.html

Shows how the multi-processing iterator creates worker processes, each with its own index_queue, and starts them with the _worker_loop target. The index_queue uses cancel_join_thread to avoid deadlocks.

```python
# torch.utils.data.dataLoader.py
class _MultiProcessingDataLoaderIter(_BaseDataLoaderIter):
    def __init__(self, loader):
        ...
        for i in range(self._num_workers):
            # No certainty which module multiprocessing_context is
            index_queue = multiprocessing_context.Queue()  # type: ignore[var-annotated]
            # Need to `cancel_join_thread` here!
            # See sections (2) and (3b) above.
            index_queue.cancel_join_thread()
            w = multiprocessing_context.Process(
                target=_utils.worker._worker_loop,
                args=(self._dataset_kind, self._dataset, index_queue,
                        self._worker_result_queue, self._workers_done_event,
                        self._auto_collation, self._collate_fn, self._drop_last,
                        self._base_seed, self._worker_init_fn, i, self._num_workers,
                        self._persistent_workers, self._shared_seed))
            w.daemon = True
            w.start()
            self._index_queues.append(index_queue)
            self._workers.append(w)
        ...

```

--------------------------------

### Import libraries and set up device for Qwen3-4B fine-tuning

Source: https://infrasys-ai.github.io/aiinfra-docs/04Train05FineTune/Code01Qwen3SFT.html

Imports the libraries used for model loading, tokenization, training, and PEFT methods (LoRA, Prompt Tuning). Sets the compute device to CUDA if available. This setup is required before any of the fine-tuning experiments in the section.

```python
from unsloth import FastLanguageModel
from transformers import AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model, PromptTuningConfig, TaskType
from datasets import load_dataset
import torch

# 设置设备

```

--------------------------------

### pretrain_gpt.py 入口与 training.py 中 get_model 调用

Source: https://infrasys-ai.github.io/aiinfra-docs/04Train02ParallelAdv/10PPCode.html

展示 pretrain_gpt.py 的 main 函数调用 pretrain，以及 training.py 中 setup_model_and_optimizer 调用 get_model 的源码。

```python
# Megatron-LM/pretrain_gpt.py
if __name__ == "__main__":
    pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain)
    pretrain(train_valid_test_datasets_provider, model_provider,
             ModelType.encoder_or_decoder, forward_step,
             args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, store=store)

# Megatron-LM/megatron/training/training.py
model, optimizer, sched = setup_model_and_optimizer(model_provider, model_type)
model = get_model(model_provider, model_type)  # ← 这里决定 PP / VPP 形态
```

--------------------------------

### PyTorch DDP simple implementation with ResNet-50 on 2 GPUs

Source: https://infrasys-ai.github.io/aiinfra-docs/_sources/04Train01ParallelBegin/04.data_parallel_implement.md

Complete example of training ResNet-50 with PyTorch DistributedDataParallel on 2 GPUs. The first block imports the required libraries (torch, torch.distributed, torch.multiprocessing, torchvision). The second block defines the distributed environment setup (MASTER_ADDR/MASTER_PORT, dist.init_process_group), data loading with DistributedSampler, DDP wrapping, and the spawn-based entry point that launches 2 processes.

```python
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
import torchvision.datasets
import torchvision.models
import torchvision.transforms as T

```

```python
def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group('nccl', rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

transform = T.Compose([
    T.Resize(224),  
    T.ToTensor(),  
    T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])  

def demo_basic(rank, world_size):

    print(f'Running basic DDP example on rank {rank}.')
    setup(rank, world_size)

    dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset)
    train_loader = torch.utils.data.DataLoader(dataset, batch_size=1, sampler=sampler, pin_memory=True, num_workers=4)

    # create model and move it to GPU with id rank
    model = torchvision.models.resnet50(weights=True).to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    for data in train_loader:
        inputs, labels = data[0].to(rank), data[1].to(rank)
        optimizer.zero_grad()
        outputs = ddp_model(inputs)
        loss_fn(outputs, labels).backward()
        optimizer.step()

    cleanup()

def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)

if __name__ == '__main__':
    run_demo(demo_basic, 2)

```

--------------------------------

### Install dependencies and load Qwen-1.8B model

Source: https://infrasys-ai.github.io/aiinfra-docs/04Train04PostTrainRL/Code02DPOPPO.html

Installs required packages and loads the Qwen-1.8B model with tokenizer. Uses bfloat16 precision if CUDA is available. Sets the pad token to the EOS token.

```python
# 首先安装必要的依赖
%pip install torch transformers transformers_stream_generator numpy matplotlib

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np
import matplotlib.pyplot as plt
from transformers import AutoTokenizer, AutoModelForCausalLM

# 设置设备 - 优先使用 GPU 加速计算
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 加载 Qwen-1.8B 模型和分词器
model_name = "Qwen/Qwen-1_8B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token  # 设置填充标记

# 加载基础模型，使用 bfloat16 精度减少内存占用
base_model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    trust_remote_code=True,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
).to(device)
print("Qwen-1.8B 模型加载完成")
```
Tokens:
Raw