Value Residual Learning

https://github.com/zcchill/value-residual-learning
Admin
Value Residual Learning is a machine learning repository that implements Resformer and SVformer...
Tokens:5,904
Snippets:53
Trust Score:3.4
License:MIT
Update:5 months ago
Show doc for...
Tokens:
Raw
### Install Transformers Library

Source: https://github.com/zcchill/value-residual-learning/blob/main/README.md

Installs the specified version of the Transformers library. Ensure this version is used for compatibility.

```bash
pip install transformers=4.44.2
```

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

### RWKV-v7 Demo Example

Source: https://github.com/zcchill/value-residual-learning/blob/main/README.md

Illustrative code from the RWKV-LM project, specifically from the RWKV-v7 demo script.

```python
def rwkv_v7_demo(model_path):

    # Load model and tokenizer
    model = RWKV(model_path=model_path, strategy='cuda fp16')
    tokenizer = Tokenizer('rwkv_vocab_v20230424.json')

    # Generate text
    init_state = None
    for i in range(200):
        out, init_state = model.forward([tokenizer.encode('The quick brown fox')] + 
                                        ([tokenizer.emoji2id[c] for c in list(init_state['last_token_id'])] if init_state else []), 
                                        init_state)
        out = out.squeeze(0)
        out[tokenizer.eos_token_id] = -float('inf')
        out = torch.multinomial(torch.softmax(out, dim=-1), num_samples=1)
        if i % 10 == 0:
            print(tokenizer.decode(out.tolist()))
        if out.item() == tokenizer.eos_token_id:
            break

    return tokenizer.decode(out.tolist())
```

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

### U-Net Form Value Residual Example

Source: https://github.com/zcchill/value-residual-learning/blob/main/README.md

Example code snippet demonstrating a U-Net form for value residual learning, found within a modded nanogpt project.

```python
self.residual_layer = nn.Sequential(
 nn.Linear(dim, dim * 2),
 nn.GELU(),
 nn.Linear(dim * 2, dim),
 nn.Dropout(dropout),
)

def forward(self, x, residual):
 residual = self.residual_layer(residual)
 return x + residual
```

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

### x-transformers Code Example

Source: https://github.com/zcchill/value-residual-learning/blob/main/README.md

Code snippet from the x-transformers library, specifically related to transformer layers.

```python
class FeedForward(nn.Module):
    def __init__(self, dim, mult=4, dropout=0., activation=None):
        super().__init__()
        self.max_seq_len = max_seq_len
        self.dim = dim
        self.dropout = dropout
        self.activation = activation or nn.GELU()

        self.net = nn.Sequential(
            nn.Linear(dim, dim * mult), 
            self.activation,
            nn.Dropout(dropout),
            nn.Linear(dim * mult, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, **kwargs):
        return self.net(x)
```

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

### Launch Training with Accelerate

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Use this command to launch the training script with specified configurations for DeepSpeed Zero2, model parameters, dataset, and learning rate scheduling. Redirects logs to a file.

```bash
accelerate launch \
    --config_file config/accelerate_deepspeed_zero2_config.yml \
    --num_processes 8 \
    src/train.py \
    --do_train \
    --bf16 \
    --modeling_name   modeling_llama_resformer \
    --model_type      llama \
    --config_overrides "vocab_size=50277,hidden_size=512,intermediate_size=1792,num_hidden_layers=8,num_attention_heads=8,num_key_value_heads=8,max_position_embeddings=2048,hidden_act=silu,rope_theta=10000,rms_norm_eps=1e-5,attention_dropout=0.0,tie_word_embeddings=False,eos_token_id=0,bos_token_id=0,use_cache=False" \
    --tokenizer_name  data/tokenizer/RedPajama-INCITE-Base-7B \
    --attn_implementation flash_attention_2 \
    --dataset_dir     data/processed_slimpajama_20B \
    --output_dir      output/modeling_llama_resformer_82M \
    --domain_weight_train '{"RedPajamaCommonCrawl_length2048":0.5,"RedPajamaC4_length2048":0.2,"RedPajamaGithub_length2048":0.1,"RedPajamaStackExchange_length2048":0.05,"RedPajamaWikipedia_length2048":0.05,"RedPajamaBook_length2048":0.05,"RedPajamaArXiv_length2048":0.05}' \
    --domain_weight_eval  '{"RedPajamaCommonCrawl_length2048":1.0}' \
    --max_steps       10000 \
    --per_device_train_batch_size 32 \
    --gradient_accumulation_steps 4 \
    --lr_scheduler_name linear_warmup_cosine \
    --learning_rate   6e-4 \
    --lr_end          6e-5 \
    --warmup_step     120 \
    --num_warmup_stop_steps 10000 \
    --optim adamw_torch \
    --weight_decay 0.1 \
    --max_grad_norm 1.0 \
    --seed 42 \
    2>&1 | tee logs/modeling_llama_resformer_82M.log
```

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

### Create Log and Output Directories

Source: https://github.com/zcchill/value-residual-learning/blob/main/README.md

Creates necessary directories for storing logs and model outputs before training.

```bash
mkdir logs
mkdir output
```

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

### Sample N Tokens from a Domain for Training Data

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Randomly samples text from a reorganized domain directory until a target token budget is reached, then writes a single train.jsonl output file. Ensure the tokenizer path is correct.

```bash
python src_data/selection.py \
    --input_dir  data/slimpajama_all \
    --output_dir data/slimpajama_20B \
    --meta_name  RedPajamaCommonCrawl \
    --data_size  10000000000 \
    --tokenizer  data/tokenizer/RedPajama-INCITE-Base-7B \
    --cache_dir  cache/

# Output: data/slimpajama_20B/RedPajamaCommonCrawl/train.jsonl
# Run for each domain; use run_selection.sh / scp_valid.sh to batch all domains.
```

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

### Tokenize and Chunk Domain Data for Training

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Tokenizes raw text files for a single domain and writes fixed-length chunks to disk in HuggingFace Arrow format. Use the provided shell script to batch-process all domains.

```bash
python src_data/preprocess.py \
    --dataset_dir data/slimpajama_20B \
    --output_dir  data/processed_slimpajama_20B \
    --domain      RedPajamaCommonCrawl \
    --max_length  2048 \
    --nproc       8 \
    --tokenizer   data/tokenizer/RedPajama-INCITE-Base-7B \
    --cache_dir   cache/

# Or batch-process all domains:
bash src_data/run_all_preprocess.sh   # pass max_length as first argument

# Output layout:
# processed_slimpajama_20B/
#   train/
#     RedPajamaCommonCrawl_length2048/   ← Arrow shards
#     RedPajamaC4_length2048/
#     ...
#   validation/
#     RedPajamaCommonCrawl_length2048/
```

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

### Train LLaMA Resformer 82M with DeepSpeed ZeRO-2

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

This bash script initiates the training process for the LLaMA Resformer 82M model using DeepSpeed ZeRO-2 optimization across 8 GPUs. Ensure cache and code directories are correctly set in the script.

```bash
# Train LLaMA Resformer 82M on 8 GPUs with DeepSpeed ZeRO-2:
mkdir -p logs output
# Edit CACHE and CODE_DIR in the script, then:
bash scripts/run_llama_resformer_82M.sh
```

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

### Run Per-Layer Token Similarity Analysis

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Computes the average cosine similarity between token representations per layer. Edit script parameters before running. Output is a list of similarity values.

```bash
cd analyze
# Edit model_name_or_path, embedding_type ("hidden" or "key"), head_operation inside script
python get_similarity.py

# Output: list of average per-layer similarity values, e.g.
# [0.12, 0.18, 0.25, 0.31, 0.38, 0.44, 0.50, 0.57]
```

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

### Linear Warmup Exponential LR Scheduler

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Implements a learning rate scheduler with a linear warmup phase, followed by an exponential decay to a specified end learning rate, and then maintains a constant end learning rate.

```python
from trainer import LinearWarmupExponentialLR

optimizer = torch.optim.AdamW([torch.zeros(1, requires_grad=True)], lr=6e-4)
scheduler = LinearWarmupExponentialLR(
    optimizer,
    num_warmup_steps=120,
    num_warmup_stop_steps=10000,
    num_training_steps=10000,
    lr_start=1e-7,
    lr_end=6e-5,
)
```

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

### Run Llama Baseline Training

Source: https://github.com/zcchill/value-residual-learning/blob/main/README.md

Executes the training script for the Llama baseline model. Modify CACHE and CODE_DIR variables in the script before running.

```bash
bash scripts/run_llama_baseline_82M.sh
```

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

### Linear Warmup Cosine LR Scheduler

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Implements a learning rate scheduler that linearly warms up to a base learning rate, then follows a cosine decay to a specified end learning rate, and finally holds the end learning rate.

```python
import torch
from trainer import LinearWarmupCosineLR

optimizer = torch.optim.AdamW([torch.zeros(1, requires_grad=True)], lr=6e-4)
scheduler = LinearWarmupCosineLR(
    optimizer,
    num_warmup_steps=120,
    num_warmup_stop_steps=10000,
    num_training_steps=10000,
    lr_start=1e-7,
    lr_end=6e-5,
)
for step in range(10000):
    optimizer.step()
    scheduler.step()
    if step in (0, 120, 5000, 9999):
        print(f"step {step}: lr = {scheduler.get_last_lr()[0]:.2e}")
# step 0:    lr = 5.00e-07   (during warmup)
# step 120:  lr = 6.00e-04   (peak)
# step 5000: lr = ~3.06e-04  (mid cosine)
# step 9999: lr = 6.00e-05   (floor)
```

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

### Extended Training Arguments with Custom LR Schedulers

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Defines `FullTrainingArguments` which extends HuggingFace's `TrainingArguments` to include custom fields for learning rate schedulers like `lr_end` and `num_warmup_stop_steps`.

```python
from training_args import FullTrainingArguments

args = FullTrainingArguments(
    output_dir="output/my_run",
    do_train=True,
    bf16=True,
    max_steps=10000,
    per_device_train_batch_size=32,
    gradient_accumulation_steps=4,
    lr_scheduler_name="linear_warmup_cosine",  # or "linear_warmup_exponential" / "linear_warmup_linear"
    learning_rate=6e-4,
    lr_end=6e-5,          # final LR after decay  (new field)
    warmup_steps=120,
    num_warmup_stop_steps=10000,  # step at which decay ends (new field)
    weight_decay=0.1,
    optim="adamw_torch",
    adam_beta1=0.9,
    adam_beta2=0.95,
)
```

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

### Run Llama Resformer Training

Source: https://github.com/zcchill/value-residual-learning/blob/main/README.md

Executes the training script for the Llama Resformer model. Modify CACHE and CODE_DIR variables in the script before running.

```bash
bash scripts/run_llama_resformer_82M.sh
```

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

### Load Weighted Multi-Domain Preprocessed Dataset

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Loads a single IterableDataset that streams preprocessed shards from multiple domains, interleaved according to specified sampling probabilities. This is used directly by the training script.

```python
from dataloaders import get_preprocessed_mixed_dataset

domain_weight_train = {
    "RedPajamaCommonCrawl_length2048": 0.50,
    "RedPajamaC4_length2048":          0.20,
    "RedPajamaGithub_length2048":      0.10,
    "RedPajamaStackExchange_length2048": 0.05,
    "RedPajamaWikipedia_length2048":   0.05,
    "RedPajamaBook_length2048":        0.05,
    "RedPajamaArXiv_length2048":       0.05,
}

train_dataset = get_preprocessed_mixed_dataset(
    preprocessed_dir="data/processed_slimpajama_20B",
    dataset_name="slimpajama",
    domain_weight=domain_weight_train,
    split="train",
    max_samples=10_240_000,   # total examples cap
    seed=42,
    no_interleave=False,      # True → sequential domain order
    shuffle=True,
    keep_in_memory=False,
)

# Iterate
for batch in train_dataset.take(3):
    print(batch.keys())   # dict_keys(['input_ids', 'attention_mask', 'domain_id'])
```

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

### SVformer: Shared Value Projection in LLaMA Attention

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

This code illustrates the SVformer architecture where only the first layer has a `v_proj` weight, and subsequent layers reuse these value states directly. This significantly reduces the number of value projection matrices.

```python
# Key forward logic (modeling_llama_svformer.py)
# __init__: v_proj is only created when layer_idx == 0
if self.layer_idx == 0:
    self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)

# forward:
if self.layer_idx == 0:
    value_states = self.v_proj(hidden_states)
    formal_layer_values.append(value_states)
else:
    value_states = formal_layer_values[0]   # direct reuse, no blending

# Train SVformer (82 M):
bash scripts/run_llama_svformer_82M.sh
```

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

### Build a HuggingFace Trainer-compatible data collator

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Use `get_data_collator` to create a collator function for HuggingFace's Trainer. It can optionally pad sequences and always sets `labels = input_ids` with padding tokens masked.

```python
from transformers import AutoTokenizer
from dataloaders import get_data_collator

tokenizer = AutoTokenizer.from_pretrained("data/tokenizer/RedPajama-INCITE-Base-7B", use_fast=True)

collator = get_data_collator(
    tokenizer,
    do_padding=False,   # sequences are already fixed-length after preprocessing
    max_length=2048,
)

# collator receives a list of dicts with 'input_ids', 'attention_mask'
batch = collator([
    {"input_ids": [1, 2, 3, 4], "attention_mask": [1, 1, 1, 1]},
    {"input_ids": [5, 6, 7, 8], "attention_mask": [1, 1, 1, 1]},
])
# batch["input_ids"]      → torch.LongTensor (2, 4)
# batch["labels"]         → torch.LongTensor (2, 4), same as input_ids
# batch["attention_mask"] → torch.LongTensor (2, 4)
```

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

### Re-organize SlimPajama Data by Domain

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

This script reorganizes the raw SlimPajama-627B dataset into separate JSONL files per domain and chunk. Set Data_Path and Write_Dir at the top of the file before running.

```python
# reorganize_data.py (run as a script; set Data_Path and Write_Dir at top of file)
# Data_Path = "Value-Residual-Learning/data/SlimPajama-627B"
# Write_Dir = "Value-Residual-Learning/data/slimpajama_all/{}"

# Produces files like:
#   slimpajama_all/RedPajamaCommonCrawl/train_chunk1.jsonl
#   slimpajama_all/RedPajamaC4/train_chunk1.jsonl
#   ...
#   slimpajama_all/RedPajamaWikipedia/validation_chunk_all.jsonl

python src_data/reorganize_data.py
```

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

### Analyze Per-Layer Attention Entropy

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

A Python script to calculate the Shannon entropy of attention distributions for each layer of a LLaMA model. Requires a JSON list of text samples and model path to be set within the script.

```bash
# Place sampled_data.json in analyze/ (a JSON list of strings, e.g. ["text1", "text2", ...])
# Set model_name_or_path inside the script, then:
cd analyze
python get_entropy.py

# Output (printed to stdout): list of average per-layer entropy values, e.g.
# [2.31, 2.45, 2.67, 2.89, 3.01, 2.95, 2.78, 2.60]   (one float per layer)
```

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

### Plot Relative Loss Analysis

Source: https://github.com/zcchill/value-residual-learning/blob/main/README.md

Runs the Python script to generate plots for relative loss analysis.

```bash
analyze/plot_relative_loss.py
```

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

### Plot Relative Training Loss

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Plots the difference in training loss between a method and a baseline. Configure log files and run the script. Output is a PDF figure.

```python
# Configure file_ls_all at the bottom of plot_relative_loss.py:
file_ls_all = {
    "ResformerVsBaseline": [[
        ["logs/modeling_llama_82M.log",        "logs/modeling_llama_resformer_82M.log"],
        ["logs/modeling_llama_82M.log",        "logs/modeling_llama_svformer_82M.log"],
    ], [
        "Resformer",
        "SVformer",
    ]],
}

# Then run:
python analyze/plot_relative_loss.py
# Output: figure_v2/ResformerVsBaseline.pdf
# X-axis: Training Step (0–10000)
# Y-axis: Relative Training Loss (method loss − baseline loss)
# Dashed red line at y=0 marks the baseline
```

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

### Calculate Token Entropy

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Calculates the entropy of token importance based on attention weights. Requires scipy and numpy.

```python
from scipy.stats import entropy as scipy_entropy
import numpy as np, torch

def calculate_entropy(attention_weights):
    # attention_weights: (batch=1, num_heads, seq, seq)
    average_attention = torch.mean(attention_weights, dim=1).squeeze(0)  # (seq, seq)
    token_importance  = average_attention.sum(dim=0)                      # (seq,)
    token_importance  = np.array(token_importance.tolist())
    token_importance  = token_importance / token_importance.sum()
    return scipy_entropy(token_importance)
```

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

### Resformer: Value Residual Blending in LLaMA Attention

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

This code snippet shows the key forward logic for Resformer, where each LLaMA attention layer blends its value projection with the first layer's value states. Formal layer values are accumulated and passed through the decoder stack.

```python
# Key forward logic inside LlamaAttention (modeling_llama_resformer.py)
value_states = self.v_proj(hidden_states)
if self.layer_idx == 0:
    formal_layer_values.append(value_states)   # store layer-0 values
else:
    value_states = 0.5 * formal_layer_values[0] + 0.5 * value_states  # residual blend

# Instantiate and train via train.py:
from modeling.modeling_llama_resformer import LlamaForCausalLM
from transformers import LlamaConfig

config = LlamaConfig(
    vocab_size=50277, hidden_size=512, intermediate_size=1792,
    num_hidden_layers=8, num_attention_heads=8, num_key_value_heads=8,
    max_position_embeddings=2048, hidden_act="silu",
)
model = LlamaForCausalLM(config)
print(sum(p.numel() for p in model.parameters()) / 1e6, "M params")  # ~82 M
```

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

### Calculate Token Representation Similarity

Source: https://context7.com/zcchill/value-residual-learning/llms.txt

Calculates the average cosine similarity between token representations. Supports averaging or concatenating attention heads. Normalizes data before computing similarity.

```python
def calculate_similarity(data, head_operation='average'):
    # data: (1, num_heads, seq, dim) for key; (1, seq, dim) for hidden
    if head_operation == 'average':
        data = data.mean(dim=1)        # (1, seq, dim)
    elif head_operation == 'concatenate':
        data = data.transpose(1, 2).reshape(data.size(0), data.size(2), -1)
    data = data.squeeze()              # (seq, dim)
    data_norm = torch.nn.functional.normalize(data, p=2, dim=-1)
    sim_matrix = torch.matmul(data_norm, data_norm.T)   # (seq, seq)
    mask = torch.eye(sim_matrix.size(0), device=data.device).bool()
    return sim_matrix[~mask].mean().item()
```
Tokens:
Raw