Chain-of-Experts

https://github.com/zihanwang314/coe
Admin
Chain-of-Experts (CoE) is a novel architecture for Mixture-of-Experts (MoE) models that enables...
Tokens:5,028
Snippets:18
Trust Score:8.4
Update:1 year ago
Show doc for...
Tokens:
Raw
### Setup Environment with Shell Script

Source: https://github.com/zihanwang314/coe/blob/main/README.MD

This script is used to set up the necessary environment for the project. It handles dependencies and configurations required to run the experiments.

```bash
bash scripts/setup.sh
```

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

### Run Basic Experiments with Shell Script

Source: https://github.com/zihanwang314/coe/blob/main/README.MD

This script executes the core experiments for evaluating the Chain-of-Experts model. It's a convenient way to start the benchmark runs.

```bash
bash runs/run_latest.sh
```

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

### Train Model with Custom Trainer Class and Distributed Setup (Python)

Source: https://context7.com/zihanwang314/coe/llms.txt

Illustrates how to initialize and use a custom `BaseTrainer` for model training. This snippet sets up distributed training using `torch.distributed.device_mesh` and FSDP, with optional sequence parallelism. It configures the trainer with a `train.yaml` configuration file and initiates the training process using `trainer.fit()`. The expected output shows training progress, validation metrics, and checkpointing information.

```python
# coe/trainer/base_trainer.py
from coe.trainer.base_trainer import BaseTrainer
from torch.distributed.device_mesh import init_device_mesh
from verl.utils.distributed import initialize_global_process_group
import hydra

@hydra.main(config_path='config', config_name='train', version_base=None)
def train_model(config):
    # Initialize distributed training
    local_rank, rank, world_size = initialize_global_process_group()

    # Create device mesh for FSDP
    device_mesh = init_device_mesh(
        device_type='cuda',
        mesh_shape=(world_size,),
        mesh_dim_names=('fsdp',)
    )

    # Create Ulysses mesh for sequence parallelism (optional)
    dp_size = world_size // config.ulysses_sequence_parallel_size
    ulysses_device_mesh = init_device_mesh(
        device_type='cuda',
        mesh_shape=(dp_size, config.ulysses_sequence_parallel_size),
        mesh_dim_names=('dp', 'sp')
    )

    # Initialize trainer
    trainer = BaseTrainer(
        config=config,
        device_mesh=device_mesh,
        ulysses_device_mesh=ulysses_device_mesh
    )

    # Start training
    trainer.fit()

if __name__ == '__main__':
    train_model()

# Expected output during training:
# Epoch 1/10 [Step 1/1000]: loss=2.45, lr=0.00003, grad_norm=1.52
# Epoch 1/10 [Step 100/1000]: loss=1.23, lr=0.0003, grad_norm=0.89
# Validation at step 100: val/loss=1.18
# Memory allocated: 45.3 GB, reserved: 48.2 GB
# Checkpoint saved: output/global_step_1000
```

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

### Load Multiple Parquet Files into Dataset (Python)

Source: https://context7.com/zihanwang314/coe/llms.txt

This example shows how to load data from multiple Parquet files into a single `BaseDataset`. It aggregates records from different sources, allowing for a combined training or validation set. The output indicates the number of records loaded from each file and the total combined count.

```python
from coe.utils.dataset.base_dataset import BaseDataset

# Combine multiple data sources
dataset = BaseDataset(
    parquet_files=[
        'data/metamathqa/train.parquet',
        'data/gsm8k/train.parquet',
        'data/additional_math/train.parquet'
    ],
    tokenizer='gpt2',
    text_keys=[['problem', 'solution']],
    max_length=1024,
    truncation='right'
)

print(f"Total samples: {len(dataset)}")

# Expected output:
# Loading Parquet files...
# Loaded 50000 records (file 1)
# Loaded 7473 records (file 2)
# Loaded 25000 records (file 3)
# Combined: 82473 total records
# Processing text fields...
# Total samples: 82473
```

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

### Python Expert Utilization Statistics Tracking

Source: https://context7.com/zihanwang314/coe/llms.txt

This Python script, using the `collections.defaultdict` from the Python standard library, is designed to track expert utilization statistics. It initializes a defaultdict to store counts for each expert, likely to be populated by processing model outputs or logs. This snippet is a starting point for gathering metrics on how often each expert is activated.

```python
import torch
from collections import defaultdict
```

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

### CoE Variant with Outer Residual Connections (Mathematical Formulation)

Source: https://github.com/zihanwang314/coe/blob/main/report.MD

This mathematical formulation details a variant of the CoE architecture employing outer residual connections. In this setup, the residual connection is applied after the intra-layer iterative processing, and the gating mechanism is dependent on the time step.

```math
x^{(0)} = x

x^{(t)} = \sum_{i=1}^{K/C} g_{t,i} \cdot \text{FFN}_i(x^{(t-1)}), \quad t = 1, 2, ..., C

y = x^{(C)} + x
```

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

### Train CoE Model with Custom Configuration (Bash)

Source: https://context7.com/zihanwang314/coe/llms.txt

This script demonstrates how to set up the environment and launch the training of a CoE model using `torchrun`. It specifies key parameters such as the number of experts, iterations, batch sizes, and data paths. The expected output includes training loss, validation loss, and checkpoint saving information.

```bash
# Set up the environment
export PYTHONPATH=/path/to/coe:$PYTHONPATH
export CUDA_VISIBLE_DEVICES="0,1,2,3"
export MASTER_PORT=29501

# Train CoE with 2 iterations, 4 experts per token, 64 total experts
torchrun --master_port=$MASTER_PORT --nproc_per_node=4 main.py \
    model.override_config.num_experts_per_tok=4 \
    model.override_config.inner_iter=2 \
    model.override_config.n_routed_experts=63 \
    model.override_config.n_shared_experts=1 \
    model.override_config.inner_residual=true \
    model.override_config.outer_residual=false \
    model.override_config.use_igate=true \
    trainer.total_training_steps=10000 \
    trainer.experiment_name=coe-64ept-2iter-4topk \
    trainer.validation_interval_steps=100 \
    trainer.save_interval_steps=1000 \
    trainer.project_name=coe-experiments \
    trainer.logger='["console","wandb"]' \
    data.train_files=data/metamathqa/train.parquet \
    data.val_files=data/metamathqa/test.parquet \
    data.max_length=512 \
    data.train_batch_size=256 \
    data.micro_batch_size_per_gpu=64

# Expected output:
# Total training steps: 10000
# MODEL TOTAL PARAMS: 544000000
# Using sequence parallel size: 1
# Step 100/10000: train/loss=1.25, train/lr=0.00028, val/loss=1.20
# Step 1000/10000: train/loss=1.15, val/loss=1.12
# Checkpoint saved to: ./outputs/coe-64ept-2iter-4topk/global_step_1000
```

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

### Save and Load HuggingFace Model Checkpoints with FSDP

Source: https://context7.com/zihanwang314/coe/llms.txt

Demonstrates how to save a model checkpoint using PyTorch FSDP's state dict capabilities, ensuring proper saving of model weights and tokenizer configurations. It also shows how to load a saved checkpoint for inference using HuggingFace's AutoModel and AutoTokenizer, and then generate text.

```python
# Save checkpoint during training
from torch.distributed.fsdp import FullStateDictConfig, StateDictType
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Assuming self.fsdp_model, self.config, and self.device_mesh are available within a class context
# self.fsdp_model: The FSDP-wrapped model
# self.config: Configuration object containing trainer.default_local_dir
# self.device_mesh: PyTorch distributed device mesh object

def save_checkpoint(self, step):
    cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
    with FSDP.state_dict_type(self.fsdp_model, StateDictType.FULL_STATE_DICT, cfg):
        state_dict = self.fsdp_model.state_dict()

    path = f"{self.config.trainer.default_local_dir}/global_step_{step}"

    if self.device_mesh.get_rank() == 0:
        # Assuming self.model is the base HuggingFace model and self.tokenizer is the tokenizer
        self.model.save_pretrained(path, state_dict=state_dict)
        self.tokenizer.save_pretrained(path)
        print(f"Checkpoint saved: {path}")

    torch.distributed.barrier() # Ensure all ranks synchronize before proceeding

# --- Example Usage for Loading and Inference ---

# Load checkpoint for inference
model_path = "output/global_step_10000"
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto" # Automatically maps layers to available devices
)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# Generate text
prompt = "Question: What is 25 * 16? Answer:"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda") # Assuming CUDA is available
outputs = model.generate(**inputs, max_length=100, temperature=0.7)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

```

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

### Run Detailed Experiments with Shell Script

Source: https://github.com/zihanwang314/coe/blob/main/README.MD

This script is used to run the detailed experiments, likely involving specific configurations and parameters for evaluating CoE's performance and advantages.

```bash
bash runs/run.sh
```

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

### Configure CoE Model Architecture (Python)

Source: https://context7.com/zihanwang314/coe/llms.txt

This Python code defines the configuration for a CoE model using HuggingFace's `PretrainedConfig`. It sets standard transformer parameters along with CoE-specific settings like the number of experts, iterations, and gating mechanisms. The configuration can be saved to a file for later use.

```python
# config/models/coe_deepseekv2/configuration_coe.py
from transformers.configuration_utils import PretrainedConfig

config = CoeConfig(
    vocab_size=102400,
    hidden_size=1024,
    num_hidden_layers=4,
    num_attention_heads=8,

    # CoE-specific parameters
    n_routed_experts=63,          # Number of routed experts per layer
    n_shared_experts=1,            # Always-active shared experts
    num_experts_per_tok=4,         # Experts selected per iteration
    inner_iter=2,                  # Number of CoE iterations

    # Architecture choices
    inner_residual=True,           # Residual connection after each iteration
    outer_residual=False,          # No residual around entire iteration block
    use_igate=True,                # Independent gating per iteration

    # Expert configuration
    moe_intermediate_size=704,     # Expert hidden dimension
    expert_hidden_size=2048,       # Expert projection size
    moe_layer_freq=1,              # Every layer is MoE (=1)

    # Training parameters
    aux_loss_alpha=0.001,          # Load balancing loss weight
    scoring_func='softmax',        # Expert weight computation
    norm_topk_prob=False,          # Don't normalize top-k weights

    max_position_embeddings=2048,
    rope_theta=10000.0
)

# Save configuration
config.save_pretrained("./config/models/my_coe_model")

# Expected: Configuration files created in ./config/models/my_coe_model/
```

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

### Python Training Loop in BaseTrainer.fit()

Source: https://context7.com/zihanwang314/coe/llms.txt

This Python code defines the main training loop within the BaseTrainer.fit() method. It handles data loading, iteration, moving data to the GPU, executing training and validation steps, and saving checkpoints. It iterates until a total number of steps is reached, reshuffling data if the iterator is exhausted.

```python
def fit(self):
    rank = self.device_mesh.get_rank()
    global_step = 0

    # Training loop
    train_iterator = iter(self.train_dataloader)
    while global_step < self.total_steps:
        try:
            data = next(train_iterator)
        except StopIteration:
            self.train_sampler.set_epoch(epoch=epoch + 1)
            train_iterator = iter(self.train_dataloader)
            data = next(train_iterator)

        # Convert to TensorDict and move to GPU
        data = TensorDict(data, batch_size=self.config.data.train_batch_size).cuda()

        # Training step
        metric = self.training_step(data)

        if rank == 0:
            print(f"Step {global_step}: {metric}")

        global_step += 1

        # Validation
        if global_step % self.config.trainer.validation_interval_steps == 0:
            val_losses = []
            for val_data in self.val_dataloader:
                val_data = TensorDict(
                    val_data,
                    batch_size=self.config.data.micro_batch_size_per_gpu
                ).cuda()
                val_loss = self.validation_step(val_data)
                val_losses.append(val_loss)

            avg_val_loss = torch.mean(torch.stack(val_losses))
            print(f"Validation loss: {avg_val_loss.item():.4f}")

        # Save checkpoint
        if global_step % self.config.trainer.save_interval_steps == 0:
            self.save_checkpoint(step=global_step)
```

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

### Create Custom Dataset with Tokenization (Python)

Source: https://context7.com/zihanwang314/coe/llms.txt

Demonstrates how to create a custom dataset using `BaseDataset` from parquet files. It tokenizes text using `AutoTokenizer` and handles multiple text keys, specifying maximum length and truncation strategy. The output shows the shapes of processed tensors, including input IDs, attention mask, position IDs, and loss mask.

```python
# coe/utils/dataset/base_dataset.py
from coe.utils.dataset.base_dataset import BaseDataset
from transformers import AutoTokenizer

# Initialize tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Create dataset with multiple text keys
dataset = BaseDataset(
    parquet_files=['data/metamathqa/train.parquet'],
    tokenizer=tokenizer,
    text_keys=[['query', 'response']],  # Combines query and response
    max_length=512,
    truncation='right'  # Options: 'right', 'left', 'error'
)

# Access dataset items
sample = dataset[0]
print(f"Input IDs shape: {sample['input_ids'].shape}")
print(f"Attention mask shape: {sample['attention_mask'].shape}")
print(f"Position IDs shape: {sample['position_ids'].shape}")
print(f"Loss mask shape: {sample['loss_mask'].shape}")

# Expected output:
# Loaded 395000 records
# Processed 395000 texts
# Input IDs shape: torch.Size([512])
# Attention mask shape: torch.Size([512])
# Position IDs shape: torch.Size([512])
# Loss mask shape: torch.Size([512])

# Each item returns a dictionary with:
# - input_ids: Token IDs (padded/truncated to max_length)
# - attention_mask: 1 for real tokens, 0 for padding
# - position_ids: Position indices for rotary embeddings
# - loss_mask: Mask for computing loss (same as attention_mask)
```

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

### Track expert usage across layers and iterations

Source: https://context7.com/zihanwang314/coe/llms.txt

This Python script iterates through samples, encodes them, and collects expert indices for each layer and iteration. It then analyzes the load balancing across experts by calculating the max, min, and standard deviation of expert counts for each iteration per layer.

```python
from collections import defaultdict
import torch
import numpy as np

# Assuming tokenizer, texts, model, and config are already defined and loaded
# tokenizer = ...
# texts = [...] # List of input texts
# model = ... # Your CoE model
# config = ... # Your model configuration

expert_counts = defaultdict(lambda: defaultdict(int))

for sample_idx in range(1000):
    input_ids = tokenizer.encode(texts[sample_idx], return_tensors="pt").to("cuda")

    with torch.no_grad():
        _ = model(input_ids)

    # Collect expert indices for each layer
    for layer in range(config.num_hidden_layers):
        for iteration in range(2):
            # This assumes expert_indices are saved to disk for each layer and iteration
            # You might need to adjust the path or how you access these indices based on your setup
            try:
                expert_indices = torch.load(
                    f"outputs/routing_logits/layer_{layer}/iter_{iteration}_topk_idx.pt"
                )

                for expert_id in expert_indices.cpu().numpy():
                    expert_counts[layer][(iteration, expert_id)] += 1
            except FileNotFoundError:
                print(f"Warning: File not found for layer {layer}, iteration {iteration}. Skipping.")
                continue # Skip to next iteration if file not found

# Analyze load balancing
for layer in range(config.num_hidden_layers):
    # Ensure we have counts for all possible experts (0 to 62, assuming 63 experts)
    # If an expert was never chosen, its count will be 0 due to defaultdict(int)
    iter0_counts = [expert_counts[layer][(0, i)] for i in range(63)]
    iter1_counts = [expert_counts[layer][(1, i)] for i in range(63)]

    print(f"Layer {layer}:")
    print(f"  Iter 0 - Max: {max(iter0_counts)}, Min: {min(iter0_counts)}, "
          f"Std: {np.std(iter0_counts):.2f}")
    print(f"  Iter 1 - Max: {max(iter1_counts)}, Min: {min(iter1_counts)}, "
          f"Std: {np.std(iter1_counts):.2f}")

```

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

### Python Expert Routing Analysis with Transformers

Source: https://context7.com/zihanwang314/coe/llms.txt

This Python script analyzes expert routing patterns in a causal language model, specifically for the Chain-of-Experts (COE) architecture. It loads a pre-trained model, enables routing logit saving, processes test data, and visualizes the routing decisions between experts across selected layers using matplotlib. Dependencies include PyTorch and Hugging Face Transformers.

```python
# scripts/analyze_routing.py
import torch
import matplotlib.pyplot as plt
from transformers import AutoTokenizer, AutoModelForCausalLM
from config.models.coe_deepseekv2.modeling_coe import CoeForCausalLM
import numpy as np

# Load model with routing logit saving enabled
model_name = "chain-of-experts/64ept-4tpk-2itr-metamathqa-10k"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# Load and configure model
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,
    trust_remote_code=True
)

config = model.config
config.save_routing_logits = True
model = CoeForCausalLM(config)
model.load_state_dict(torch.load(f"{model_name}/pytorch_model.bin"))
model.eval()
model.to("cuda")

# Analyze routing on test data
selected_layers = [0, 1, config.num_hidden_layers - 2, config.num_hidden_layers - 1]
routing_matrix = {layer: torch.zeros(63, 63) for layer in selected_layers}

for idx in range(1000):
    text = "question: What is 2+2? answer: 4"
    input_ids = tokenizer.encode(text, return_tensors="pt").to("cuda")

    with torch.no_grad():
        _ = model(input_ids)

    # Load routing decisions for each layer and iteration
    for layer in selected_layers:
        iter0_idx = torch.load(f"outputs/routing_logits/layer_{layer}/iter_0_topk_idx.pt")
        iter1_idx = torch.load(f"outputs/routing_logits/layer_{layer}/iter_1_topk_idx.pt")

        # Count expert transitions between iterations
        for expert0, expert1 in zip(iter0_idx.cpu(), iter1_idx.cpu()):
            routing_matrix[layer][expert0, expert1] += 1

# Visualize routing patterns
fig, axes = plt.subplots(1, len(selected_layers), figsize=(12, 3))
for i, layer in enumerate(selected_layers):
    matrix = routing_matrix[layer].numpy()
    axes[i].imshow(matrix, cmap='viridis', aspect='auto')
    axes[i].set_title(f"Layer {layer}")
    axes[i].set_xlabel("Expert in Iter 1")
    axes[i].set_ylabel("Expert in Iter 0")

plt.tight_layout()
plt.savefig("outputs/routing_analysis.png", dpi=300)
```

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

### Previous MoE Output Representation (Mathematical Formula)

Source: https://github.com/zihanwang314/coe/blob/main/report.MD

This snippet presents the mathematical formula for representing the output of a previous Mixture-of-Experts (MoE) layer, where experts process tokens independently and in parallel. It defines the gate values and expert outputs.

```mathematica
y = \sum_{i=1}^K g_i \cdot E_i(x)

g_i = \begin{cases}
 s_i, & s_i \in \text{TopK}(s_j|1 \leq j \leq K, N) \\
 0, & \text{otherwise}
\end{cases}

s_i = \text{Softmax}(u_i^\top e_i)
```

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

### CoE Iterative Processing Mechanism (Mathematical Formula)

Source: https://github.com/zihanwang314/coe/blob/main/report.MD

This snippet defines the iterative processing mechanism for Chain-of-Experts (CoE). It shows how the output of one iteration ($x^{(t-1)}$) becomes the input for the next iteration ($x^{(t)}$), incorporating expert computations and residual connections.

```mathematica
x^{(0)} = x

x^{(t)} = \sum_{i=1}^{K/C} g_{t,i} \cdot \text{FFN}_i(x^{(t-1)}) + \mathbb{I}_r \cdot x^{(t-1)}, \quad t = 1, 2, ..., C

y = x^{(C)}
```

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

### CoE Variant with Shared Gating (Mathematical Formulation)

Source: https://github.com/zihanwang314/coe/blob/main/report.MD

This mathematical formulation describes a variant of the Chain-of-Experts (CoE) architecture that utilizes a shared gating mechanism instead of independent gating. It shows how the input is processed iteratively through experts and then combined.

```math
x^{(0)} = x

x^{(t)} = \sum_{i=1}^{K/C} g_i \cdot \text{FFN}_i(x^{(t-1)}) + I_r \cdot x^{(t-1)}, \quad t = 1, 2, ..., C

y = x^{(C)}
```

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

### CoE Independent Gating Mechanism (Mathematical Formula)

Source: https://github.com/zihanwang314/coe/blob/main/report.MD

This snippet details the independent gating mechanism used in Chain-of-Experts (CoE), adapted from DeepSeek-V2. It describes how gate values are determined for each iteration and expert, utilizing different routers and token embeddings.

```mathematica
g_{t,i} = \begin{cases}
 s_{t,i}, & s_{t,i} \in \text{TopK}(s_{t,j}|1 \leq j \leq K/C, N) \\
 0, & \text{otherwise}
\end{cases}

s_{t,i} = \text{Softmax}(u_{t,i}^\top e_{t,i})
```
Tokens:
Raw