orbit-torch 0.0.4a1__py3-none-any.whl → 0.1.0b1__py3-none-any.whl

orbit_torch-0.1.0b1.dist-info/METADATA

@@ -0,0 +1,208 @@
+Metadata-Version: 2.4
+Name: orbit-torch
+Version: 0.1.0b1
+Summary: A PyTorch training engine with plugin system and advanced model components
+Home-page: https://github.com/A03HCY/Orbit
+Author: Aiden Hopkins
+Author-email: acdphc@qq.com
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0.0
+Requires-Dist: rich
+Requires-Dist: tensorboard
+Requires-Dist: matplotlib
+Requires-Dist: seaborn
+Requires-Dist: numpy
+Requires-Dist: scikit-learn
+Requires-Dist: einops
+Requires-Dist: tokenizers
+Requires-Dist: transformers
+Requires-Dist: safetensors
+Requires-Dist: accelerate
+Requires-Dist: lpips
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Orbit
+
+Orbit is a flexible, plugin-based PyTorch training engine designed to simplify the training loop while providing powerful components for modern deep learning models, including LLMs.
+
+It features a modular design with a rich set of plugins, advanced model building blocks (like MoE, RoPE, GQA), comprehensive LoRA/DoRA support, and cutting-edge optimizers.
+
+## Features
+
+### 🚀 Core Engine
+- **Plugin System**: Decoupled training logic using plugins for callbacks, logging, and training strategies.
+- **Simplified Loop**: Clean `train` and `eval` interfaces.
+- **Flexible Updates**:
+  - `auto_update()`: Automatically handles forward pass, loss calculation, backward pass, optimizer step, and zero grad.
+  - `update(loss)`: Allows manual control over the update step if you need custom forward/loss logic.
+
+### 🧩 Model Components (`orbit.model`)
+Orbit provides a collection of high-performance, reusable layers:
+- **Attention**: `MultiHeadAttention` with support for **GQA** (Grouped Query Attention), **RoPE** (Rotary Positional Embeddings), and FlashAttention.
+- **LoRA & DoRA**: Full support for Low-Rank Adaptation and **Weight-Decomposed Low-Rank Adaptation (DoRA)** across `Linear`, `Conv2d`, `Conv1d`, and `Embedding` layers. Also supports **Gated LoRA**.
+- **MoE**: Mixture of Experts block with `TopKGate` routing.
+- **Gates**: A variety of gating mechanisms including `SigmoidGate`, `TanhGate`, `SoftmaxGate`, `GLUGate`, `ContextGate`, and `TopKGate`.
+- **Others**: `FiLM` (Feature-wise Linear Modulation), `MLP` (with Gated support), `RotaryPositionalEmbedding`.
+
+### 🛠️ Utilities & Kit (`orbit.utils`)
+
+Orbit provides a comprehensive toolkit to speed up development:
+
+#### 🔧 LoRA Utilities
+Manual control over LoRA injection and management (alternative to the Plugin approach).
+- **Injection**:
+  - `inject_lora(model, r=8, ...)`: Manually inject LoRA/DoRA/Gated LoRA into specific layers.
+  - `inject_lora_file(model, path)`: Automatically inject and load LoRA configuration/weights from a file.
+- **Management**:
+  - `merge_lora(model)` / `unmerge_lora(model)`: Merge weights for faster inference or unmerge to resume training.
+  - `save_lora(model, path)` / `load_lora(model, path)`: Efficiently save/load only LoRA parameters.
+  - `freeze_backbone_only(model)`: Helper to freeze the base model while keeping LoRA and specified heads trainable.
+- **Diagnosis**:
+  - `LoRADiagnoser`: Check for rank collapse and monitor gradient norms during training.
+
+#### ❄️ Model Freezing
+- `freeze_layers(model, targets=['encoder'])`: Freeze layers matching the target names (supports wildcards).
+- `unfreeze_layers(model, targets)`: Unfreeze specific layers.
+- `get_trainable_params(model)`: Get parameters for the optimizer.
+
+#### 🎭 Masking
+- `make_causal_mask`: Create causal masks for autoregressive models.
+- `make_padding_mask`, `make_lookahead_mask`, `make_sliding_window_mask`.
+
+#### 💾 Layer I/O
+- `save_layer(model, layer_name, path)`: Save weights of a specific sub-module (e.g., just the backbone).
+- `load_layer(model, layer_name, path)`: Load weights into a specific sub-module.
+- `get_model_by_name(model, name)`: Access sub-modules using dot notation strings (e.g., "backbone.layer1").
+
+#### 📝 SFT Helpers
+- `build_sft`: Prepares data for Supervised Fine-Tuning (handles chat templates, tokenization, and label masking).
+- `train_sft(engine)`: A specialized training step for SFT that handles the forward pass and loss calculation automatically.
+
+#### ⚙️ Optimization (`orbit.optim`)
+- **Muon**: MomentUm Orthogonalized by Newton-schulz optimizer.
+- **SAM**: Sharpness-Aware Minimization wrapper.
+
+#### 🌱 Initialization & Seeding
+- `auto_initialize(model)`: Automatically initializes weights based on layer type (Linear, Conv, Embedding, etc.).
+- `seed_everything(seed)`: Sets seeds for Python, NumPy, PyTorch, and CUDA for reproducibility.
+
+#### 🖥️ CUDA
+- `cuda_alloc(size)`: Optimizes PyTorch CUDA memory allocation configuration (e.g., `max_split_size_mb`).
+
+### 🔌 Plugins (`orbit.plugin`)
+- `EarlyStopping`: Stop training when a metric stops improving.
+- `GradientAccumulation`: Simulate larger batch sizes.
+- `Warmup`: Learning rate warmup.
+- `Mentor`: Training assistant/logger.
+- `MemoryEstimator`: Monitor CUDA memory usage.
+- `LoRA`: Easy injection of LoRA layers via plugin.
+- `Board`: TensorBoard integration.
+
+## Installation
+
+```bash
+pip install orbit-torch
+```
+
+**Requirements**:
+- Python >= 3.8
+- PyTorch >= 2.0.0 (Required for FlashAttention backend)
+
+## Quick Start
+
+### 1. Basic Training (CIFAR-10)
+
+```python
+import torch
+import torch.nn as nn
+from orbit.engine import Engine
+from orbit.plugin import EarlyStopping, GradientAccumulation, Mentor
+from orbit.utils import auto_initialize
+
+# Define your model
+model = MyConvNet()
+auto_initialize(model)
+
+# Setup Engine
+trainer = Engine(
+    model=model,
+    criterion=nn.CrossEntropyLoss(),
+    optimizer=torch.optim.Adam(model.parameters(), lr=1e-3),
+    plugins=[
+        Mentor(),
+        EarlyStopping(monitor='val_acc', patience=3),
+        GradientAccumulation(steps=2)
+    ]
+)
+
+# Train
+for _ in trainer.train(train_loader, num_epochs=10):
+    trainer.auto_update() # Handles forward, backward, step, zero_grad
+    
+    # Handle Epoch End (e.g., Validation)
+    if not trainer.is_epoch_end: continue
+    
+    for _ in trainer.eval(test_loader): 
+        trainer.auto_update()
+```
+
+### 2. LLM SFT with LoRA/DoRA
+
+Orbit makes it easy to fine-tune LLMs using LoRA or DoRA.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from orbit.engine import Engine
+from orbit.plugin import LoRA, GradientAccumulation
+from orbit.utils import train_sft, seed_everything
+
+seed_everything(42)
+
+# Load Model
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B", device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
+
+# Setup Engine with LoRA Plugin
+trainer = Engine(
+    model=model,
+    optimizer=torch.optim.AdamW(model.parameters(), lr=1e-4),
+    plugins=[
+        # Inject DoRA into MLP layers
+        LoRA(target_names=['mlp'], dora=True, r=16, alpha=32),
+        GradientAccumulation(steps=8)
+    ]
+)
+
+# Train Loop
+# Assuming `dataloader` yields SFT batches (input_ids, attention_mask, labels)
+for _ in trainer.train(dataloader, num_epochs=3):
+    # train_sft handles the forward pass and loss calculation for CausalLM
+    train_sft(trainer) 
+```
+
+### 3. Chat Interface
+
+Interact with your trained model in the terminal:
+
+```python
+from orbit.kit import ChatInterface
+
+chat = ChatInterface(model_id="path/to/model", device="cuda")
+chat.interact()
+```
+
+## License
+
+MIT License

orbit_torch-0.1.0b1.dist-info/RECORD

@@ -0,0 +1,65 @@
+orbit/__init__.py,sha256=Pnt9BItQOlPZR4StcWQ3C-XmYHvqmLSyJiF0honrEew,78
+orbit/callback.py,sha256=qI4iaHwY5AXpzgSHNdiAEgK5XJxv8eEMTg9YxpzVeFk,1669
+orbit/engine.py,sha256=-S0kNNu_rgX17xrurKDgigqjzgiYN3AKoD4QDXqi80g,35894
+orbit/dataset/__init__.py,sha256=OAhQmm3GYwOfsznpGRNEhPr0Lanw_iIrqJc7FV2gbBc,46
+orbit/dataset/cogn.py,sha256=kN_CPkV_VnsS7pVN0tdXsah0dYr780K7gfo3ROIGWyY,4749
+orbit/dataset/data/cogn_en.jsonl,sha256=Ba_KQKW3BJs3jWD7B_z0_Lwq34Zujgu0M2xe7RWO0dA,14862
+orbit/dataset/data/cogn_zh.jsonl,sha256=d__Dr_5Z3w33KcX8CikK34VNt9DVAFUnDvnhCJikumk,39055
+orbit/kit/__init__.py,sha256=40DZDZwGKgOeJ0vUsqhCBKHt-H1Tj0a79dYrjfKKOgs,82
+orbit/kit/interface.py,sha256=_eB6PguJk-EjpLjnCuGDTHJ2woEhHxqW36sasCYC9Ws,5658
+orbit/kit/wrapper.py,sha256=PSBDEx4fsnuznZ3-9EYcUTr4hgQsRPDT9EiI-rxvcaY,5941
+orbit/model/__init__.py,sha256=DY9le8idTuOgCu-K2Os4j83Fa6X55G3VT1DpGP6GvTs,151
+orbit/model/base.py,sha256=NekTUd1P_0e8yHsRUoIJWJbNYqf0-XdDTgFk9WeNBOQ,4264
+orbit/model/config.py,sha256=2IMe6bxES_qf1hXNhroMuAosN9BCZfJJTMJryrz0Q2c,1888
+orbit/model/registry.py,sha256=r1nMCYTpBT24Fu-mOL-S37zOQA-ngpd4uHvonnkN6Lo,1606
+orbit/model/block/__init__.py,sha256=NvNs7F1MfUxOmpHQgrTE0XcJ20OOreTBCGKTYxtggtc,938
+orbit/model/block/attention.py,sha256=7eSE0GM5wW2gsjcTdqUoHEyd8VqCJ1bNbfF-HY7mAUU,10131
+orbit/model/block/bio.py,sha256=SNCG1urEfqDoX-B7-hNpHWJ_vHrU8EBJP9-sjdeh2Nw,21291
+orbit/model/block/codebook.py,sha256=IhxOp0_--euNvrNCU95_h8e38XA6HBuZef8dDC7Ex-g,4195
+orbit/model/block/conv.py,sha256=NswgOt5daz60ROOqmorrftcxI5I0aDQ14rnorYfUcpM,18616
+orbit/model/block/embedding.py,sha256=hMFsx2gYhm_qNoF22Xg2rkOKf5IP16fflUttV8q1OA8,9101
+orbit/model/block/film.py,sha256=oBWmm2mclSjE4yyVf-Qy8n5AeoGkj4xBGERylqNFkMs,6542
+orbit/model/block/fusion.py,sha256=YbyFiIOgbYD0ywrqIPPr0T662yEmc4VcJBT0OSR-yEw,12361
+orbit/model/block/gate.py,sha256=ODE-bKRVaqkTzhAI3z_JCNINfCq1fosQHSrWHqeTmBk,10648
+orbit/model/block/lora.py,sha256=ggmbj9_qQu0FzkqE1A1XRnuIv_bOCVDIPergN-kZPQc,31208
+orbit/model/block/mlp.py,sha256=EcJWbLREPoBVCgQDHW6nJ-2zuxZ0J-biFhZNXd0BFRQ,2226
+orbit/model/block/moe.py,sha256=JW-1esQ986T6_wDQwF7065RjkTSqPx9vEcNppSPG0tE,3151
+orbit/model/block/tcn.py,sha256=ClkYw2cMFcXFm6mMbrUZnraZ5YtokGotU1_sJBGANAU,3836
+orbit/model/kit/__init__.py,sha256=-Q8mGi39sDmUZP9P3nZ1GoyvEOz-rfa0_JEXSzlfplg,156
+orbit/model/kit/discriminator.py,sha256=JA30tHCjmZ1L1oPf-7o5w1nM-_9h1gB9PiPkkNz_JSE,1484
+orbit/model/kit/losses.py,sha256=CSR3aH6SeUIA72sopi2XAQDDHyFelKhZAZcNG3tLKmM,7697
+orbit/model/motif/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+orbit/model/motif/vision/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+orbit/model/motif/vision/v1.py,sha256=a1guncpocl1x57pGpVa60FR_7rHVBV3v_JOVLgLee4o,22981
+orbit/optim/__init__.py,sha256=tO24DdBQ7hfBbk2tKIwh5Y0WXqEYjxowLuwXCuScFMI,45
+orbit/optim/muon.py,sha256=IT0l1b0mAtLpB-SUuPyo9uaYHO2x_wkqg1TLcDJ6Bvk,7342
+orbit/optim/sam.py,sha256=aeqWF7aCfZuhuzlSXFEGP8fUTbEDp4Cm75Q-Xr8I6Ho,3949
+orbit/plugin/__init__.py,sha256=VJbpVtYaGB8roeSrv1PCaGsKmlqDTHpJ91LkeEh1xHg,466
+orbit/plugin/board.py,sha256=ofuxEJf1CC2I7Y5fR_3_Ln04HPQqBKTVp2-8CU7lvI8,2266
+orbit/plugin/checkpoint.py,sha256=R-R-180KWcF264DPvWgjGAoWtX-Vkn0nS0MQjOBfDkU,15264
+orbit/plugin/classification.py,sha256=h2I_a19NzC8RLrO5qD1TbaZ6QBh3Wig8m0-kBiXNdn4,7020
+orbit/plugin/display_model.py,sha256=DuvefqcHCHUCyS8bybBzxaXr6Y8gxlpql-Sa9hePBuw,2906
+orbit/plugin/early_stopping.py,sha256=E8l4vSiZ1oECK9SeXFcrhMuDEC72PGHTwB4p_oZghIc,4060
+orbit/plugin/ema.py,sha256=4HnVOfFDRwdjUeOPkkOVj-A6fA8GPEBetsFPlI5wVuk,3887
+orbit/plugin/gradient_accumulation.py,sha256=Hr2TyPDUx8SOBJ6IpBfP5wPBvKj-JgYBNlF28LIs9-U,1061
+orbit/plugin/lora.py,sha256=LmpBc1Fo54eANsAVf36W93O8tLiKcv9z1Ofv8PgtIJ0,17032
+orbit/plugin/memory_estimator.py,sha256=7YrlAJb2yOh6nTaNK29NAu_aLWHpboIN8mnqYIL4LZo,9738
+orbit/plugin/mentor.py,sha256=sx3gLHnxYJX1tJhqnqHPTW2KuTeNM7h0PqOIquxMaxQ,12474
+orbit/plugin/overfit.py,sha256=lpfxX4igi4uNdY4OR2_fpayw52iEfkCQnQ0HgzodZd4,1194
+orbit/plugin/warmup.py,sha256=XvU47kNSk43OSVPkS1riyhndySOK-Vo5M_SReUjxSHo,4991
+orbit/plugin/data/mentor_i18n.json,sha256=KHnEKRKHeLga_nqEUZVgiTh4jb4j2f_yD0MJT-7sdis,8005
+orbit/utils/__init__.py,sha256=2rxWZ27U4oZH0XxqicYOV8oeWQh2NTv-igakgdRPbMI,1093
+orbit/utils/cuda.py,sha256=5JUxXVDJ_IzQfEI9oguTKSz2uSLT2GTFNXbxSZVWhFM,267
+orbit/utils/freeze.py,sha256=ee6MxqoLaOZ-kf5DZyzUAsS5JUCFQPU_2LxkLZQIxJo,4265
+orbit/utils/image.py,sha256=wklEHDcH_TGvT3aDBPKSJvJWHGgHu4yxssuVfZ8u4No,6775
+orbit/utils/initialization.py,sha256=tIRQMlyNfcVDvrtLwLKthpyCjJ_Cy7AN5ZkggFt9o0s,23407
+orbit/utils/layer_io.py,sha256=_qGmt2GDOxVb0RacTK4ZMALek36oQLYVqCeuwcXzaZE,3920
+orbit/utils/lora.py,sha256=LXqSuh9QGTRn_WJsOnJPmiSy0uzpFnWpLxAmLWq_twc,18458
+orbit/utils/mask.py,sha256=DayMxmKlWMYMYPceYhYlf9lIx4FD9s4Tp1WvzfVpsu0,3072
+orbit/utils/moe.py,sha256=miLesFrUHIe9gWr0STs4d4NgnUbV7nskq5N4ywRtEj0,2132
+orbit/utils/seed.py,sha256=PaApRMEO2codOXDcmelybRJ2ZgAk-Sh4K5SHlZ4KpHE,1789
+orbit/utils/sft.py,sha256=Y6hDGW4tAgtzmEeSqolkYcBKYrOJvuW1vxgJMWvzdRU,3127
+orbit_torch-0.1.0b1.dist-info/METADATA,sha256=IzDxu8AeHbXbrub3QisQybVGdYoYzjx3c1t12MRIzZE,7831
+orbit_torch-0.1.0b1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+orbit_torch-0.1.0b1.dist-info/top_level.txt,sha256=emrF0of931NzTSL4R5yBKpGoewFCB-cAwYNcUF5cqBs,6
+orbit_torch-0.1.0b1.dist-info/RECORD,,

orbit_torch-0.0.4a1.dist-info/METADATA

@@ -1,25 +0,0 @@
-Metadata-Version: 2.4
-Name: orbit-torch
-Version: 0.0.4a1
-Summary: A PyTorch training engine with plugin system
-Author: Aiden Hopkins
-Author-email: acdphc@qq.com
-Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: OS Independent
-Requires-Python: >=3.8
-Requires-Dist: torch>=1.10.0
-Requires-Dist: rich
-Requires-Dist: tensorboard
-Requires-Dist: matplotlib
-Requires-Dist: seaborn
-Requires-Dist: numpy
-Requires-Dist: scikit-learn
-Requires-Dist: einops
-Requires-Dist: tokenizers
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary

orbit_torch-0.0.4a1.dist-info/RECORD

@@ -1,29 +0,0 @@
-orbit/__init__.py,sha256=xikJCuBzK0VSwJ8gvPt0cZijw5TzVI2t3a62_aK-uEo,51
-orbit/callback.py,sha256=FXq-bOVfoYW0S5S95ry55yCn1QgYkecpHonO35BNKwE,1738
-orbit/engine.py,sha256=4Vskd39cfVa7jUWucikva8oDc68om4DtGH6OZLUDhKE,34230
-orbit/optim/__init__.py,sha256=jp3TZFLM5LRcHeeYjR6qNwzwo1dfeR65A6DSGCEPnFg,67
-orbit/optim/muon.py,sha256=IT0l1b0mAtLpB-SUuPyo9uaYHO2x_wkqg1TLcDJ6Bvk,7342
-orbit/optim/sam.py,sha256=lDfou1jQYM_qtlWgYzfUKp2Zdi-uBSKZS7iYxB-yGkU,3702
-orbit/plugin/__init__.py,sha256=9EcMoS__kzldM4MBNUwv5xUs21rliqrBALVmzIk_ADw,435
-orbit/plugin/board.py,sha256=dK3aaTmznq9vSYvpSPgfE9kt30tieR93oK0oDaC8yJw,2267
-orbit/plugin/checkpoint.py,sha256=w9R5O2-jzbRPSGPiuGuUycq21EebAdqigEFK0T0waOE,10898
-orbit/plugin/classification.py,sha256=iVOFNUaL4SUMNEtnmbnioprb2nH1sOV4IsGH7UPhU8U,7019
-orbit/plugin/display_model.py,sha256=Uch563Jq0R78lGHYiDPotXB6dmWfqdClSOiuOIKw7i8,2907
-orbit/plugin/early_stopping.py,sha256=eDTjIzxDZvnCXANR2RXn8vrMh5wX8F2rQjl0BNp-6dU,4065
-orbit/plugin/ema.py,sha256=NtQVB3Rz5YreQwWicXCNy_CfD6d7pu0crGGrz6-IJI8,3892
-orbit/plugin/gradient_accumulation.py,sha256=JaSl-U9Zo_jeEaJE-KAXEPhQdjFA_He729VaCw2W8n4,1066
-orbit/plugin/memory_estimator.py,sha256=YxSXbjAwl6jJh6yFEod5Tu_pgZh1Qadb0JsgmyfF7Mc,9743
-orbit/plugin/mentor.py,sha256=sx3gLHnxYJX1tJhqnqHPTW2KuTeNM7h0PqOIquxMaxQ,12474
-orbit/plugin/overfit.py,sha256=lpfxX4igi4uNdY4OR2_fpayw52iEfkCQnQ0HgzodZd4,1194
-orbit/plugin/warmup.py,sha256=jS3Q52zuhTdqN9Bw85zcoH1YXZoqyRvnD7JHSSLRUwQ,4996
-orbit/plugin/data/mentor_i18n.json,sha256=KHnEKRKHeLga_nqEUZVgiTh4jb4j2f_yD0MJT-7sdis,8005
-orbit/utils/__init__.py,sha256=gXNfTz1a_kL5UZXT2o5BhoItc4Dgrk8w4K2hsfkbiYk,589
-orbit/utils/freeze.py,sha256=ujR5wl2GHuqIQi0rm2i_yz6W8dj6X4YqgyZ3tBbv4yU,2405
-orbit/utils/initialization.py,sha256=X0-OP0FPWX60_SRuB-x_Od3KdDh4EAD9UbOCZFj6NtI,20560
-orbit/utils/layer_io.py,sha256=E-YnYw7av2ZdczDzhs08hUQtbZrcHPvqDdzEvWgs7Gs,1872
-orbit/utils/mask.py,sha256=DayMxmKlWMYMYPceYhYlf9lIx4FD9s4Tp1WvzfVpsu0,3072
-orbit/utils/seed.py,sha256=tbjF2jIPQfb0M_-wXOMVH3_XzWVsDkuhzD5lctGqalA,2578
-orbit_torch-0.0.4a1.dist-info/METADATA,sha256=g2NSH_PVm9afSyVAfvct5xvk2hFColq85tKdbGQfNFc,700
-orbit_torch-0.0.4a1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-orbit_torch-0.0.4a1.dist-info/top_level.txt,sha256=emrF0of931NzTSL4R5yBKpGoewFCB-cAwYNcUF5cqBs,6
-orbit_torch-0.0.4a1.dist-info/RECORD,,