WCNegentropy
/

BitTransformerLM

@@ -1,485 +0,0 @@
-#!/usr/bin/env python3
-"""
-BitTransformerLM TRUE 1.21B Parameter Training
-==============================================
-The REAL DEAL: 1.21B parameters with PROPER FSDP sharding (not duplication!)
-Based on our proven 680M success, now scaled to the full billion+ parameters!
-"""
-import os
-import sys
-import time
-import json
-import logging
-import argparse
-import torch.multiprocessing as mp
-from datetime import datetime
-from typing import Dict, Any, Optional
-import torch
-import torch.nn as nn
-import torch.distributed as dist
-import torch.nn.functional as F
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy
-from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
-from torch.utils.data import DataLoader, DistributedSampler
-from datasets import load_dataset
-from bit_transformer.model import BitTransformerLM
-from bit_transformer.bit_io import text_to_bits
-from bit_transformer.utils import set_dropout
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
-logger = logging.getLogger(__name__)
-class True1BConfig:
-    """TRUE 1.21B parameter configuration with optimized settings."""
-    # Model Architecture - FULL 1.21B parameters
-    D_MODEL = 2048
-    NUM_LAYERS = 24
-    NUM_HEADS = 32
-    DIM_FEEDFORWARD = 8192
-    MAX_SEQ_LEN = 512  # Optimized length from our 680M success
-    # Training Configuration
-    BATCH_SIZE_PER_GPU = 1  # Conservative
-    NUM_GPUS = 4
-    GRADIENT_ACCUMULATION_STEPS = 32
-    EFFECTIVE_BATCH_SIZE = BATCH_SIZE_PER_GPU * NUM_GPUS * GRADIENT_ACCUMULATION_STEPS  # 128
-    LEARNING_RATE = 2e-4
-    WEIGHT_DECAY = 0.01
-    MAX_STEPS = 1000  # Reasonable for demo
-    WARMUP_STEPS = 100
-    # OPTIMIZED BitTransformerLM settings (proven to work)
-    USE_REVERSIBLE = True
-    USE_GRADIENT_CHECKPOINTING = True
-    USE_MIXED_PRECISION = True
-    CHUNK_SIZE = 128  # Chunked attention for memory efficiency
-    FULL_ATTN_LOGGING = False  # Memory optimization
-    # Reduced telemetry impact (proven necessary)
-    LAMBDA_K = 0.1
-    LAMBDA_C = 0.1
-    LAMBDA_S = 0.1
-    @classmethod
-    def get_model_config(cls) -> Dict[str, Any]:
-        """Get optimized model configuration."""
-        return {
-            "d_model": cls.D_MODEL,
-            "nhead": cls.NUM_HEADS,
-            "num_layers": cls.NUM_LAYERS,
-            "dim_feedforward": cls.DIM_FEEDFORWARD,
-            "max_seq_len": cls.MAX_SEQ_LEN,
-            "lambda_K": cls.LAMBDA_K,
-            "lambda_C": cls.LAMBDA_C,
-            "lambda_S": cls.LAMBDA_S,
-            "reversible": cls.USE_REVERSIBLE,
-            "use_checkpoint": cls.USE_GRADIENT_CHECKPOINTING,
-            "use_autocast": True,
-            "chunk_size": cls.CHUNK_SIZE,
-            "full_attn_logging": cls.FULL_ATTN_LOGGING,
-        }
-class OptimizedWikiTextDataset(torch.utils.data.Dataset):
-    """Optimized WikiText dataset for 1.21B training."""
-    def __init__(self, split: str = "train", max_samples: int = 1000, max_length: int = 512):
-        self.max_length = max_length
-        logger.info(f"Loading WikiText-103 {split} (max {max_samples} samples)...")
-        dataset = load_dataset("wikitext", "wikitext-103-raw-v1", split=split)
-        # Get good samples
-        texts = [item['text'] for item in dataset
-                if len(item['text'].strip()) > 50][:max_samples]
-        self.texts = texts
-        logger.info(f"Loaded {len(self.texts)} samples from {split}")
-    def __len__(self) -> int:
-        return len(self.texts)
-    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
-        text = self.texts[idx]
-        try:
-            bits = text_to_bits(text)
-            if len(bits) > self.max_length:
-                bits = bits[:self.max_length]
-            elif len(bits) < self.max_length:
-                bits = bits + [0] * (self.max_length - len(bits))
-            input_bits = torch.tensor(bits[:-1], dtype=torch.long)
-            target_bits = torch.tensor(bits[1:], dtype=torch.long)
-            return {
-                'input_ids': input_bits,
-                'labels': target_bits
-            }
-        except Exception:
-            # Fallback pattern
-            pattern = [0, 1] * (self.max_length // 2)
-            input_bits = torch.tensor(pattern[:-1], dtype=torch.long)
-            target_bits = torch.tensor(pattern[1:], dtype=torch.long)
-            return {
-                'input_ids': input_bits,
-                'labels': target_bits
-            }
-def setup_distributed(rank: int, world_size: int) -> None:
-    """Setup distributed training."""
-    os.environ['MASTER_ADDR'] = 'localhost'
-    os.environ['MASTER_PORT'] = '29500'
-    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
-    dist.init_process_group("nccl", rank=rank, world_size=world_size)
-    torch.cuda.set_device(rank)
-def cleanup_distributed() -> None:
-    """Cleanup distributed training."""
-    dist.destroy_process_group()
-def create_fsdp_model(config: True1BConfig, rank: int) -> FSDP:
-    """Create PROPERLY SHARDED FSDP model (not duplicated!)."""
-    logger.info("🏗️ Creating TRUE 1.21B parameter model with PROPER FSDP sharding...")
-    model_config = config.get_model_config()
-    # Create model on CPU first
-    model = BitTransformerLM(**model_config)
-    params = sum(p.numel() for p in model.parameters())
-    if rank == 0:
-        logger.info(f"✅ Base model: {params:,} parameters ({params/1e9:.2f}B)")
-    # PROPER FSDP configuration for SHARDING (not duplication)
-    fsdp_config = {
-        "auto_wrap_policy": size_based_auto_wrap_policy,
-        "sharding_strategy": ShardingStrategy.FULL_SHARD,  # FULL SHARDING!
-        "mixed_precision": MixedPrecision(
-            param_dtype=torch.float16,
-            reduce_dtype=torch.float16,
-            buffer_dtype=torch.float16,
-        ),
-        "backward_prefetch": BackwardPrefetch.BACKWARD_PRE,
-        "device_id": rank,
-        "limit_all_gathers": True,
-        "use_orig_params": False,  # Memory optimization
-    }
-    # Wrap with FSDP for SHARDING
-    model = FSDP(model, **fsdp_config)
-    if rank == 0:
-        logger.info("✅ FSDP model created with FULL SHARDING (not duplication)")
-        logger.info("🚀 Each GPU handles 1/4 of the 1.21B parameters!")
-    return model
-def train_step(model: FSDP, batch: Dict[str, torch.Tensor],
-               optimizer: torch.optim.Optimizer, scaler: torch.cuda.amp.GradScaler,
-               rank: int) -> tuple:
-    """Optimized training step."""
-    model.train()
-    input_ids = batch['input_ids'].to(rank, non_blocking=True)
-    labels = batch['labels'].to(rank, non_blocking=True)
-    with torch.cuda.amp.autocast():
-        outputs = model(input_ids)
-        if isinstance(outputs, tuple):
-            logits, telemetry = outputs
-        else:
-            logits, telemetry = outputs, {}
-        loss = F.cross_entropy(logits.view(-1, 2), labels.view(-1))
-    scaler.scale(loss).backward()
-    return loss.item(), telemetry
-def save_checkpoint(model: FSDP, optimizer, scheduler, step: int,
-                   config: True1BConfig, rank: int) -> str:
-    """Save 1.21B parameter checkpoint."""
-    if rank == 0:
-        checkpoint_dir = f"/data/checkpoints/true_1b_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
-        os.makedirs(checkpoint_dir, exist_ok=True)
-        # Save FSDP state dict
-        with FSDP.state_dict_type(model, FSDP.StateDictType.FULL_STATE_DICT):
-            model_state = model.state_dict()
-        checkpoint = {
-            'step': step,
-            'model_state_dict': model_state,
-            'optimizer_state_dict': optimizer.state_dict(),
-            'scheduler_state_dict': scheduler.state_dict(),
-            'config': config.get_model_config(),
-            'timestamp': datetime.now().isoformat(),
-            'parameters': 1210000000,  # Approximate
-        }
-        checkpoint_path = f"{checkpoint_dir}/model.pt"
-        torch.save(checkpoint, checkpoint_path)
-        logger.info(f"💾 1.21B model saved: {checkpoint_path}")
-        return checkpoint_path
-    return ""
-def test_inference(model: FSDP, config: True1BConfig, rank: int) -> Dict[str, Any]:
-    """Test inference with the trained 1.21B model."""
-    if rank != 0:
-        return {}
-    logger.info("🧪 Testing 1.21B parameter model inference...")
-    model.eval()
-    set_dropout(model, 0.0)
-    inference_results = []
-    # Test patterns
-    test_patterns = [
-        "Hello world",
-        "The quick brown fox",
-        "In the beginning",
-        "Once upon a time",
-        "Artificial intelligence"
-    ]
-    with torch.no_grad():
-        for i, text in enumerate(test_patterns):
-            try:
-                # Convert to bits
-                bits = text_to_bits(text)
-                if len(bits) > config.MAX_SEQ_LEN - 50:  # Leave room for generation
-                    bits = bits[:config.MAX_SEQ_LEN - 50]
-                input_bits = torch.tensor(bits, dtype=torch.long).unsqueeze(0).to(rank)
-                # Generate continuation
-                with torch.cuda.amp.autocast():
-                    for _ in range(20):  # Generate 20 more bits
-                        outputs = model(input_bits)
-                        if isinstance(outputs, tuple):
-                            logits, telemetry = outputs
-                        else:
-                            logits = outputs
-                            telemetry = {}
-                        # Get next bit prediction
-                        next_bit_logits = logits[0, -1, :]
-                        next_bit = torch.softmax(next_bit_logits, dim=-1).argmax().item()
-                        # Append to sequence
-                        next_tensor = torch.tensor([[next_bit]], dtype=torch.long).to(rank)
-                        input_bits = torch.cat([input_bits, next_tensor], dim=1)
-                        if input_bits.size(1) >= config.MAX_SEQ_LEN:
-                            break
-                # Convert back to text
-                generated_bits = input_bits.squeeze().cpu().tolist()
-                try:
-                    generated_text = bits_to_text(generated_bits)
-                except:
-                    generated_text = f"[Generated {len(generated_bits)} bits]"
-                result = {
-                    'input': text,
-                    'input_bits': len(bits),
-                    'generated_bits': len(generated_bits),
-                    'output': generated_text[:200],  # Limit length
-                    'telemetry': {k: float(v) if isinstance(v, torch.Tensor) else v
-                                for k, v in telemetry.items()}
-                }
-                inference_results.append(result)
-                logger.info(f"Test {i+1}: '{text}' -> Generated {len(generated_bits)} bits")
-            except Exception as e:
-                logger.warning(f"Inference test {i+1} failed: {e}")
-                inference_results.append({
-                    'input': text,
-                    'error': str(e)
-                })
-    logger.info("✅ 1.21B model inference testing complete!")
-    return {'inference_results': inference_results}
-def main_worker(rank: int, world_size: int, config: True1BConfig) -> None:
-    """Main training worker for 1.21B model."""
-    setup_distributed(rank, world_size)
-    if rank == 0:
-        logger.info("🚀 TRUE 1.21B PARAMETER BITTRANSFORMERLM TRAINING!")
-        logger.info("=" * 60)
-        logger.info("✅ PROPER FSDP SHARDING (not duplication)")
-        logger.info("✅ Based on proven 680M success")
-        logger.info("✅ All optimizations enabled")
-    # Create datasets
-    train_dataset = OptimizedWikiTextDataset("train", max_samples=2000, max_length=config.MAX_SEQ_LEN)
-    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
-    train_loader = DataLoader(
-        train_dataset,
-        batch_size=config.BATCH_SIZE_PER_GPU,
-        sampler=train_sampler,
-        num_workers=0,  # Avoid multiprocessing issues
-        pin_memory=True
-    )
-    # Create FSDP model with PROPER sharding
-    model = create_fsdp_model(config, rank)
-    # Setup optimizer and scheduler
-    optimizer = torch.optim.AdamW(
-        model.parameters(),
-        lr=config.LEARNING_RATE,
-        weight_decay=config.WEIGHT_DECAY,
-        betas=(0.9, 0.95)
-    )
-    scheduler = torch.optim.lr_scheduler.OneCycleLR(
-        optimizer,
-        max_lr=config.LEARNING_RATE,
-        total_steps=config.MAX_STEPS,
-        pct_start=config.WARMUP_STEPS / config.MAX_STEPS,
-    )
-    scaler = torch.cuda.amp.GradScaler()
-    if rank == 0:
-        logger.info("🎯 Starting 1.21B parameter training...")
-    # Training loop
-    step = 0
-    running_loss = 0.0
-    start_time = time.time()
-    checkpoint_path = ""
-    try:
-        for epoch in range(10):
-            train_sampler.set_epoch(epoch)
-            for batch_idx, batch in enumerate(train_loader):
-                loss, telemetry = train_step(model, batch, optimizer, scaler, rank)
-                running_loss += loss
-                # Gradient accumulation
-                if (batch_idx + 1) % config.GRADIENT_ACCUMULATION_STEPS == 0:
-                    scaler.unscale_(optimizer)
-                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
-                    scaler.step(optimizer)
-                    scaler.update()
-                    scheduler.step()
-                    optimizer.zero_grad()
-                    step += 1
-                    # Logging
-                    if step % 10 == 0 and rank == 0:
-                        avg_loss = running_loss / 10
-                        elapsed = time.time() - start_time
-                        memory_used = torch.cuda.memory_allocated(rank) / (1024**3)
-                        logger.info(
-                            f"Step {step:4d} | "
-                            f"Loss: {avg_loss:.4f} | "
-                            f"K: {telemetry.get('negentropy', 0):.3f} | "
-                            f"C: {telemetry.get('lz_complexity', 0):.3f} | "
-                            f"S: {telemetry.get('symbiosis', 0):.3f} | "
-                            f"LR: {scheduler.get_last_lr()[0]:.2e} | "
-                            f"Mem: {memory_used:.1f}GB | "
-                            f"Time: {elapsed:.1f}s"
-                        )
-                        running_loss = 0.0
-                        start_time = time.time()
-                    # Save checkpoint
-                    if step % 100 == 0 and step > 0:
-                        checkpoint_path = save_checkpoint(model, optimizer, scheduler, step, config, rank)
-                    if step >= config.MAX_STEPS:
-                        break
-            if step >= config.MAX_STEPS:
-                break
-        # Final checkpoint
-        if rank == 0:
-            checkpoint_path = save_checkpoint(model, optimizer, scheduler, step, config, rank)
-            logger.info("🏆 1.21B PARAMETER TRAINING COMPLETED SUCCESSFULLY!")
-            # Test inference
-            inference_results = test_inference(model, config, rank)
-            # Save results to benchmarks
-            benchmark_data = {
-                'timestamp': datetime.now().isoformat(),
-                'model_parameters': '1.21B',
-                'training_steps': step,
-                'final_loss': running_loss,
-                'checkpoint_path': checkpoint_path,
-                'inference_results': inference_results,
-                'config': config.get_model_config(),
-            }
-            with open('/data/true_1b_results.json', 'w') as f:
-                json.dump(benchmark_data, f, indent=2)
-            logger.info("📊 Results saved to /data/true_1b_results.json")
-    except Exception as e:
-        if rank == 0:
-            logger.error(f"Training failed: {e}")
-        raise
-    finally:
-        cleanup_distributed()
-def main():
-    """Main entry point."""
-    config = True1BConfig()
-    world_size = 4
-    if not torch.cuda.is_available() or torch.cuda.device_count() < world_size:
-        print("❌ Need 4 CUDA GPUs for 1.21B training!")
-        return
-    print("🚀 Launching TRUE 1.21B parameter training with PROPER FSDP sharding!")
-    print("🎯 This will work because we've proven the hardware capability!")
-    # Launch distributed training
-    mp.spawn(
-        main_worker,
-        args=(world_size, config),
-        nprocs=world_size,
-        join=True
-    )
-if __name__ == "__main__":
-    main()