🚀 Refined BitTransformerLM: Organized codebase with best practices

bit_transformer/model.py

@@ -1,20 +1,21 @@
-import math
+"""BitTransformerLM model implementation with reversible layers and telemetry."""
+
 import contextlib
 import logging
-from typing import Dict, List, Tuple, Optional
+import math
+import sys
+from typing import Dict, List, Optional, Tuple
 
 import torch
 import torch.distributed as dist
-import sys
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.checkpoint as checkpoint
 
-from .torch_utils import cpu_autocast
-
-from .optimization import configure_optimizer
 from .compression import decompress_bits
+from .optimization import configure_optimizer
 from .parity import enforce_parity
+from .torch_utils import cpu_autocast
 
 _mask_cache: Dict[Tuple[int, torch.device], torch.Tensor] = {}
 _attention_cache: Dict[str, torch.Tensor] = {}  # For caching attention patterns
@@ -29,7 +30,15 @@ def clear_cache():
 
 
 def get_tri_mask(seq_len: int, device: torch.device) -> torch.Tensor:
-    """Return or create a cached upper-triangular mask with memory management."""
+    """Return or create a cached upper-triangular mask with memory management.
+
+    Args:
+        seq_len: Sequence length for the mask.
+        device: PyTorch device for tensor allocation.
+
+    Returns:
+        Upper-triangular boolean mask tensor.
+    """
     key = (seq_len, device)
     
     # Clear cache if it gets too large
@@ -56,7 +65,12 @@ except Exception:  # pragma: no cover - handle missing torch or unsupported vers
 
 
 class PositionalEncoding(nn.Module):
-    """Sinusoidal positional encoding."""
+    """Sinusoidal positional encoding for transformer inputs.
+
+    Args:
+        d_model: Model dimension for embedding.
+        max_len: Maximum sequence length to precompute.
+    """
 
     def __init__(self, d_model: int, max_len: int = 1024) -> None:
         super().__init__()
@@ -70,7 +84,14 @@ class PositionalEncoding(nn.Module):
         self.register_buffer("pe", pe.unsqueeze(1))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Add positional encoding to input tensor."""
+        """Add positional encoding to input tensor.
+
+        Args:
+            x: Input tensor of shape (seq_len, batch_size, d_model).
+
+        Returns:
+            Input tensor with positional encoding added.
+        """
         return x + self.pe[: x.size(0)]
 
 
@@ -325,7 +346,14 @@ class ReversibleLoggingTransformerEncoderLayer(nn.Module):
 
 
 class BitTransformerLM(nn.Module):
-    """Transformer language model that operates on raw bits (0/1) with telemetry."""
+    """Bit-native transformer language model with reversible layers and telemetry.
+
+    A transformer architecture that processes binary sequences directly with:
+    - Reversible layers for memory efficiency
+    - Built-in safety telemetry (K/C/S metrics)
+    - Chunked attention for long sequences
+    - Causal and diffusion training modes
+    """
 
     def __init__(
         self,
@@ -349,10 +377,23 @@ class BitTransformerLM(nn.Module):
         """Create a BitTransformer language model.
 
         Args:
-            full_attn_logging: When ``False`` and ``chunk_size`` is
-                smaller than the sequence length, the model skips
-                reconstructing the full ``T×T`` attention matrices for
-                telemetry to reduce memory use.
+            d_model: Model dimension for embeddings and attention.
+            nhead: Number of attention heads.
+            num_layers: Number of transformer layers.
+            dim_feedforward: Dimension of feedforward networks.
+            max_seq_len: Maximum sequence length for positional encoding.
+            lambda_K: Weight for negentropy metric in telemetry.
+            lambda_C: Weight for complexity metric in telemetry.
+            lambda_S: Weight for symbiosis metric in telemetry.
+            reversible: Enable reversible layers for memory efficiency.
+            use_checkpoint: Use gradient checkpointing.
+            use_autocast: Use automatic mixed precision.
+            use_act: Enable Adaptive Computation Time.
+            act_threshold: ACT halting threshold.
+            chunk_size: Chunk size for chunked attention (None for full attention).
+            overlap: Overlap size for chunked attention.
+            full_attn_logging: When False and chunk_size is smaller than sequence
+                length, skip reconstructing full attention matrices for telemetry.
         """
         super().__init__()
         self.d_model = d_model