litetorch 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (169) hide show
  1. litetorch-0.1.0/MANIFEST.in +8 -0
  2. litetorch-0.1.0/Makefile +32 -0
  3. litetorch-0.1.0/PKG-INFO +297 -0
  4. litetorch-0.1.0/README.md +275 -0
  5. litetorch-0.1.0/include/litetorch/allocator.h +36 -0
  6. litetorch-0.1.0/include/litetorch/amp.h +116 -0
  7. litetorch-0.1.0/include/litetorch/autograd.h +87 -0
  8. litetorch-0.1.0/include/litetorch/backend.h +84 -0
  9. litetorch-0.1.0/include/litetorch/checkpoint.h +18 -0
  10. litetorch-0.1.0/include/litetorch/cl_backend.h +227 -0
  11. litetorch-0.1.0/include/litetorch/context_parallel.h +15 -0
  12. litetorch-0.1.0/include/litetorch/continuous_batching.h +47 -0
  13. litetorch-0.1.0/include/litetorch/custom_ops.h +50 -0
  14. litetorch-0.1.0/include/litetorch/data.h +73 -0
  15. litetorch-0.1.0/include/litetorch/device.h +43 -0
  16. litetorch-0.1.0/include/litetorch/device_mesh.h +27 -0
  17. litetorch-0.1.0/include/litetorch/distributed.h +342 -0
  18. litetorch-0.1.0/include/litetorch/dtensor.h +41 -0
  19. litetorch-0.1.0/include/litetorch/fsdp.h +135 -0
  20. litetorch-0.1.0/include/litetorch/fsdp_wrapper.h +18 -0
  21. litetorch-0.1.0/include/litetorch/grad_scaler.h +33 -0
  22. litetorch-0.1.0/include/litetorch/guided_decoding.h +24 -0
  23. litetorch-0.1.0/include/litetorch/jit.h +182 -0
  24. litetorch-0.1.0/include/litetorch/llm_serving.h +62 -0
  25. litetorch-0.1.0/include/litetorch/memory_manager.h +51 -0
  26. litetorch-0.1.0/include/litetorch/nn.h +357 -0
  27. litetorch-0.1.0/include/litetorch/ops.h +113 -0
  28. litetorch-0.1.0/include/litetorch/optim.h +152 -0
  29. litetorch-0.1.0/include/litetorch/platform.h +38 -0
  30. litetorch-0.1.0/include/litetorch/quantization.h +146 -0
  31. litetorch-0.1.0/include/litetorch/serialization.h +23 -0
  32. litetorch-0.1.0/include/litetorch/tensor.h +158 -0
  33. litetorch-0.1.0/include/litetorch/thread_pool.h +161 -0
  34. litetorch-0.1.0/include/litetorch/zero3_optimizer.h +33 -0
  35. litetorch-0.1.0/litetorch.egg-info/PKG-INFO +297 -0
  36. litetorch-0.1.0/litetorch.egg-info/SOURCES.txt +168 -0
  37. litetorch-0.1.0/litetorch.egg-info/dependency_links.txt +1 -0
  38. litetorch-0.1.0/litetorch.egg-info/entry_points.txt +3 -0
  39. litetorch-0.1.0/litetorch.egg-info/not-zip-safe +1 -0
  40. litetorch-0.1.0/litetorch.egg-info/top_level.txt +1 -0
  41. litetorch-0.1.0/pyproject.toml +3 -0
  42. litetorch-0.1.0/setup.cfg +4 -0
  43. litetorch-0.1.0/setup.py +88 -0
  44. litetorch-0.1.0/src/autograd/autograd_engine.cpp +126 -0
  45. litetorch-0.1.0/src/autograd/checkpoint.cpp +66 -0
  46. litetorch-0.1.0/src/autograd/saved_tensor.cpp +37 -0
  47. litetorch-0.1.0/src/autograd.cpp +1 -0
  48. litetorch-0.1.0/src/backend/backend.cpp +316 -0
  49. litetorch-0.1.0/src/backend/gpu_native/common/gpu_common.h +182 -0
  50. litetorch-0.1.0/src/backend/gpu_native/elementwise/elementwise_ops.cu +502 -0
  51. litetorch-0.1.0/src/backend/gpu_native/entrypoint.cu +388 -0
  52. litetorch-0.1.0/src/backend/gpu_native/kernels.cu +1 -0
  53. litetorch-0.1.0/src/backend/gpu_native/kernels.hip +1 -0
  54. litetorch-0.1.0/src/backend/gpu_native/math/gemm.cu +173 -0
  55. litetorch-0.1.0/src/backend/gpu_native/math/reduction.cu +44 -0
  56. litetorch-0.1.0/src/backend/gpu_native/nn/flash_attention.cu +761 -0
  57. litetorch-0.1.0/src/backend/gpu_native/nn/nn_kernels.cu +1537 -0
  58. litetorch-0.1.0/src/backend/gpu_native/optim/optimizers.cu +124 -0
  59. litetorch-0.1.0/src/bindings/python_bindings.cpp +823 -0
  60. litetorch-0.1.0/src/cl_backend/cl_backend.cpp +655 -0
  61. litetorch-0.1.0/src/cl_backend/cl_functions.cpp +45 -0
  62. litetorch-0.1.0/src/cl_backend/cl_functions.h +56 -0
  63. litetorch-0.1.0/src/cl_backend/command_graph.cpp +48 -0
  64. litetorch-0.1.0/src/cl_backend.cpp +1 -0
  65. litetorch-0.1.0/src/data/dataloader.cpp +352 -0
  66. litetorch-0.1.0/src/data/dataset.cpp +41 -0
  67. litetorch-0.1.0/src/data.cpp +1 -0
  68. litetorch-0.1.0/src/distributed/context_parallel.cpp +79 -0
  69. litetorch-0.1.0/src/distributed/distributed.cpp +1287 -0
  70. litetorch-0.1.0/src/distributed/fsdp_wrapper.cpp +159 -0
  71. litetorch-0.1.0/src/distributed/shm_ipc.cpp +385 -0
  72. litetorch-0.1.0/src/llm/continuous_batching.cpp +200 -0
  73. litetorch-0.1.0/src/llm/guided_decoding.cpp +68 -0
  74. litetorch-0.1.0/src/llm/medusa.cpp +127 -0
  75. litetorch-0.1.0/src/memory_manager/allocator.cpp +99 -0
  76. litetorch-0.1.0/src/memory_manager/memory_manager.cpp +109 -0
  77. litetorch-0.1.0/src/memory_manager.cpp +1 -0
  78. litetorch-0.1.0/src/nn/activation.cpp +36 -0
  79. litetorch-0.1.0/src/nn/attention.cpp +95 -0
  80. litetorch-0.1.0/src/nn/container.cpp +79 -0
  81. litetorch-0.1.0/src/nn/convolution.cpp +73 -0
  82. litetorch-0.1.0/src/nn/dropout.cpp +47 -0
  83. litetorch-0.1.0/src/nn/embedding.cpp +160 -0
  84. litetorch-0.1.0/src/nn/linear.cpp +51 -0
  85. litetorch-0.1.0/src/nn/llm_serving.cpp +142 -0
  86. litetorch-0.1.0/src/nn/loss.cpp +41 -0
  87. litetorch-0.1.0/src/nn/moe.cpp +281 -0
  88. litetorch-0.1.0/src/nn/nn_utils.h +46 -0
  89. litetorch-0.1.0/src/nn/normalization.cpp +60 -0
  90. litetorch-0.1.0/src/nn/parallel_linear.cpp +119 -0
  91. litetorch-0.1.0/src/nn/pooling.cpp +29 -0
  92. litetorch-0.1.0/src/nn/qlora.cpp +70 -0
  93. litetorch-0.1.0/src/nn/transformer.cpp +72 -0
  94. litetorch-0.1.0/src/nn.cpp +1 -0
  95. litetorch-0.1.0/src/ops/activation.cpp +612 -0
  96. litetorch-0.1.0/src/ops/attention.cpp +474 -0
  97. litetorch-0.1.0/src/ops/cast.cpp +48 -0
  98. litetorch-0.1.0/src/ops/checkpoint.cpp +99 -0
  99. litetorch-0.1.0/src/ops/clip_grad.cpp +56 -0
  100. litetorch-0.1.0/src/ops/convolution.cpp +579 -0
  101. litetorch-0.1.0/src/ops/custom_ops.cpp +70 -0
  102. litetorch-0.1.0/src/ops/elementwise.cpp +1632 -0
  103. litetorch-0.1.0/src/ops/flash_decoding.cpp +118 -0
  104. litetorch-0.1.0/src/ops/fused_loss.cpp +198 -0
  105. litetorch-0.1.0/src/ops/jit.cpp +539 -0
  106. litetorch-0.1.0/src/ops/kernels.cl +2034 -0
  107. litetorch-0.1.0/src/ops/kernels.cpp +1 -0
  108. litetorch-0.1.0/src/ops/linear.cpp +379 -0
  109. litetorch-0.1.0/src/ops/loss.cpp +430 -0
  110. litetorch-0.1.0/src/ops/moe_ops.cpp +196 -0
  111. litetorch-0.1.0/src/ops/normalization.cpp +666 -0
  112. litetorch-0.1.0/src/ops/pooling.cpp +470 -0
  113. litetorch-0.1.0/src/ops/quantization_ops.cpp +106 -0
  114. litetorch-0.1.0/src/ops/ring_attention.cpp +160 -0
  115. litetorch-0.1.0/src/ops/rope.cpp +155 -0
  116. litetorch-0.1.0/src/ops/scaled_matmul.cpp +46 -0
  117. litetorch-0.1.0/src/ops/w8a8_ops.cpp +110 -0
  118. litetorch-0.1.0/src/ops.cpp +1 -0
  119. litetorch-0.1.0/src/optim/adam.cpp +73 -0
  120. litetorch-0.1.0/src/optim/adamw.cpp +104 -0
  121. litetorch-0.1.0/src/optim/adamw_8bit.cpp +107 -0
  122. litetorch-0.1.0/src/optim/adamw_fp8.cpp +129 -0
  123. litetorch-0.1.0/src/optim/cosine_annealing.cpp +24 -0
  124. litetorch-0.1.0/src/optim/grad_scaler.cpp +83 -0
  125. litetorch-0.1.0/src/optim/optim_utils.h +29 -0
  126. litetorch-0.1.0/src/optim/optimizer.cpp +15 -0
  127. litetorch-0.1.0/src/optim/rmsprop.cpp +61 -0
  128. litetorch-0.1.0/src/optim/sgd.cpp +81 -0
  129. litetorch-0.1.0/src/optim/steplr.cpp +18 -0
  130. litetorch-0.1.0/src/optim/zero3_optimizer.cpp +115 -0
  131. litetorch-0.1.0/src/optim.cpp +1 -0
  132. litetorch-0.1.0/src/serialization/optimizer.cpp +395 -0
  133. litetorch-0.1.0/src/serialization/parameters.cpp +167 -0
  134. litetorch-0.1.0/src/serialization.cpp +1 -0
  135. litetorch-0.1.0/src/tensor/storage.cpp +214 -0
  136. litetorch-0.1.0/src/tensor/tensor_core.cpp +1079 -0
  137. litetorch-0.1.0/src/tensor.cpp +1 -0
  138. litetorch-0.1.0/tests/__pycache__/demo_run.cpython-312.pyc +0 -0
  139. litetorch-0.1.0/tests/__pycache__/test_litetorch.cpython-312.pyc +0 -0
  140. litetorch-0.1.0/tests/__pycache__/test_litetorch.cpython-314.pyc +0 -0
  141. litetorch-0.1.0/tests/advanced_features_test.cpp +216 -0
  142. litetorch-0.1.0/tests/advanced_optimization_stage2_test.cpp +112 -0
  143. litetorch-0.1.0/tests/advanced_serving_test.cpp +159 -0
  144. litetorch-0.1.0/tests/allocator_test.cpp +37 -0
  145. litetorch-0.1.0/tests/attention_test.cpp +77 -0
  146. litetorch-0.1.0/tests/audit_bugs_test.cpp +86 -0
  147. litetorch-0.1.0/tests/audit_optimizations_test.cpp +35 -0
  148. litetorch-0.1.0/tests/backend_detection_test.cpp +70 -0
  149. litetorch-0.1.0/tests/checkpoint_test.cpp +32 -0
  150. litetorch-0.1.0/tests/demo_run.cpp +245 -0
  151. litetorch-0.1.0/tests/demo_run.py +156 -0
  152. litetorch-0.1.0/tests/distributed_test.cpp +96 -0
  153. litetorch-0.1.0/tests/double_backward_test.cpp +35 -0
  154. litetorch-0.1.0/tests/fusion_benchmark.cpp +64 -0
  155. litetorch-0.1.0/tests/gpt_training_test.cpp +123 -0
  156. litetorch-0.1.0/tests/gpu_allocator_test.cpp +32 -0
  157. litetorch-0.1.0/tests/gpu_cast_test.cpp +52 -0
  158. litetorch-0.1.0/tests/gpu_half_perf_test.cpp +200 -0
  159. litetorch-0.1.0/tests/jit_test.cpp +102 -0
  160. litetorch-0.1.0/tests/llm_optimization_test.cpp +187 -0
  161. litetorch-0.1.0/tests/memory_leak_test.cpp +85 -0
  162. litetorch-0.1.0/tests/new_features_test.cpp +546 -0
  163. litetorch-0.1.0/tests/production_features_test.cpp +306 -0
  164. litetorch-0.1.0/tests/production_optimization_test.cpp +166 -0
  165. litetorch-0.1.0/tests/scale_upgrade_test.cpp +105 -0
  166. litetorch-0.1.0/tests/stage4_test.cpp +125 -0
  167. litetorch-0.1.0/tests/stress_test.cpp +324 -0
  168. litetorch-0.1.0/tests/test_litetorch.py +59 -0
  169. litetorch-0.1.0/tests/transformer_test.cpp +125 -0
@@ -0,0 +1,8 @@
1
+ include README.md
2
+ include LICENSE
3
+ include pyproject.toml
4
+ include setup.py
5
+ include Makefile
6
+ recursive-include include *
7
+ recursive-include src *
8
+ recursive-include tests *
@@ -0,0 +1,32 @@
1
+ CXX ?= g++
2
+ CXXFLAGS ?= -std=c++14 -O3 -fPIC -Iinclude
3
+
4
+ ifeq ($(OS),Windows_NT)
5
+ LDFLAGS ?= -shared -lpthread -lws2_32
6
+ TARGET_LIB := build/liblitetorch.so
7
+ else
8
+ LDFLAGS ?= -shared -lpthread -ldl
9
+ TARGET_LIB := build/liblitetorch.so
10
+ endif
11
+
12
+ SRC_DIR := src
13
+ BUILD_DIR := build
14
+ OBJ_DIR := $(BUILD_DIR)/objs
15
+
16
+ SRCS := $(shell find $(SRC_DIR) -name "*.cpp" -not -path "src/bindings/*")
17
+ OBJS := $(patsubst $(SRC_DIR)/%.cpp,$(OBJ_DIR)/%.o,$(SRCS))
18
+
19
+ .PHONY: all clean
20
+
21
+ all: $(TARGET_LIB)
22
+
23
+ $(TARGET_LIB): $(OBJS)
24
+ @mkdir -p $(BUILD_DIR)
25
+ $(CXX) $(LDFLAGS) $^ -o $@
26
+
27
+ $(OBJ_DIR)/%.o: $(SRC_DIR)/%.cpp
28
+ @mkdir -p $(dir $@)
29
+ $(CXX) $(CXXFLAGS) -c $< -o $@
30
+
31
+ clean:
32
+ rm -rf $(BUILD_DIR)
@@ -0,0 +1,297 @@
1
+ Metadata-Version: 2.4
2
+ Name: litetorch
3
+ Version: 0.1.0
4
+ Summary: Python bindings for LiteTorch deep learning framework
5
+ Home-page: https://github.com/nguyenminh20000/Litetorch-
6
+ Author: LiteTorch Team
7
+ Classifier: Development Status :: 4 - Beta
8
+ Classifier: Intended Audience :: Developers
9
+ Classifier: Intended Audience :: Science/Research
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: C++
12
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
13
+ Requires-Python: >=3.8
14
+ Description-Content-Type: text/markdown
15
+ Dynamic: author
16
+ Dynamic: classifier
17
+ Dynamic: description
18
+ Dynamic: description-content-type
19
+ Dynamic: home-page
20
+ Dynamic: requires-python
21
+ Dynamic: summary
22
+
23
+ > [!WARNING]
24
+ > **THIS PROJECT IS NOT YET COMPLETE AND MAY CONTAIN SOME ERRORS OR IMPROVEMENTS. WE ARE WORKING TO FIX THEM. YOU CAN ALSO CONTRIBUTE TO THE FRAMEWORK!.**
25
+
26
+
27
+ # LiteTorch framework
28
+
29
+ LiteTorch is a lightweight, high-performance deep learning framework built natively in C++14 with seamless Python bindings via `pybind11`. Designed with an intuitive PyTorch-like API, LiteTorch features dynamic autograd graph execution, memory optimization via Activation Checkpointing, distributed training primitives (FSDP, ZeRO-3), and multi-backend hardware acceleration (NVIDIA CUDA, AMD ROCm, OpenCL, and multi-threaded CPU).
30
+
31
+ ---
32
+
33
+ ## Key Features
34
+
35
+ - **PyTorch-Style Hardware Auto-Detection**:
36
+ - Automatically senses and initializes **NVIDIA CUDA** (`nvcc` + cuBLAS/cuDNN) or **AMD ROCm/HIP** (`hipcc` + rocBLAS/MIOpen) when native GPUs are present.
37
+ - Seamlessly falls back to **OpenCL** or multi-threaded **CPU** execution on systems without native GPU drivers.
38
+ - **Dual API (C++ Core & Python Bindings)**:
39
+ - High-level Python interface: `import litetorch as lt`.
40
+ - Zero performance overhead with native C++14 execution underneath.
41
+ - **Dynamic Autograd Engine**:
42
+ - Reverse-mode automatic differentiation over Directed Acyclic Graphs (DAG).
43
+ - Topological Sort DAG traversal algorithm for precise gradient accumulation.
44
+ - **Advanced Memory Management**:
45
+ - **Activation Checkpointing**: Re-computes activations during backward passes to dramatically reduce VRAM footprint.
46
+ - **LRU Storage Eviction & Caching Allocator**: Smart memory pooling and automatic LRU swap between RAM and VRAM.
47
+ - **Distributed Training Primitives**:
48
+ - **Fully Sharded Data Parallel (FSDP)** & **ZeRO-3 Optimizer**: Shards parameters, gradients, and optimizer states across GPUs.
49
+ - Inter-node communication via NCCL (NVIDIA), RCCL (AMD), Shared Memory IPC (SHM), and TCP Socket fallback.
50
+ - **Fast Build System**: Multi-core parallel Makefile (`make -j$(nproc)`) with shared library caching (`liblitetorch.so`) enabling sub-second test runs.
51
+ - **System Console Commands**: Global command registration for executing `demo_run.py` or `test_litetorch.py` directly without `./` or `python3` prefixes.
52
+
53
+ ---
54
+
55
+ ## Quick Installation & Setup
56
+
57
+ ### 1. Auto-Install C++ Build Dependencies (Linux Auto-Installer)
58
+
59
+ Automated installer script for C++ dependencies on Linux (Ubuntu, Debian, RHEL, Fedora, Arch Linux):
60
+
61
+ ```bash
62
+ ./install_deps.sh
63
+ ```
64
+
65
+ *(For detailed OS-specific C++ & GPU toolkit installation guides, see [`REQUIREMENTS_CPP.md`](file:///home/notmerblx/Pictures/Litetorch/REQUIREMENTS_CPP.md))*
66
+
67
+ ### 2. Install Python Dependencies & Package
68
+
69
+ ```bash
70
+ python3 -m pip install -r requirements.txt
71
+ python3 -m pip install -e .
72
+ ```
73
+
74
+ After installation, you can `import litetorch as lt` or run `demo_run.py` directly anywhere in your shell!
75
+
76
+ ---
77
+
78
+ ## Architecture & Core Algorithms Breakdown
79
+
80
+ ### Layered System Architecture
81
+
82
+ ```mermaid
83
+ graph TD
84
+ A["Python Layer (import litetorch as lt)"] --> B["C++ Binding Layer (pybind11)"]
85
+ B --> C["LiteTorch High-Level API (Tensor, Ops, nn::Module, optim)"]
86
+ C --> D["Autograd & Memory Engine (DAG, Checkpointing, Caching Allocator)"]
87
+ D --> E["Distributed Engine (ProcessGroup, FSDP, ZeRO-3, NCCL/RCCL)"]
88
+ E --> F1["Backend 1: Native GPU (CUDA / ROCm cuBLAS/rocBLAS)"]
89
+ E --> F2["Backend 2: OpenCL Backend"]
90
+ E --> F3["Backend 3: Multi-Threaded CPU Engine"]
91
+ ```
92
+
93
+ ---
94
+
95
+ ### 1. Dynamic Autograd Graph & Topological Sort Algorithm
96
+
97
+ Tensor operations dynamically construct a **Directed Acyclic Graph (DAG)** where each `Tensor` acts as a Node holding a weak pointer (`std::weak_ptr<Node> creator`) to the operation that produced it.
98
+
99
+ ```mermaid
100
+ graph LR
101
+ X["Tensor X (Input)"] -->|mul| H1["Tensor H1"]
102
+ X -->|mul| H1
103
+ H1 -->|add| H2["Tensor H2 (Output)"]
104
+ X -->|add| H2
105
+ ```
106
+
107
+ #### Reverse-Mode Automatic Differentiation Workflow:
108
+ 1. **Topological Sort Traversal**:
109
+ When `loss->backward()` is called, the autograd engine executes a DFS or Kahn's algorithm to sort nodes from Output (Loss) back to Inputs.
110
+ 2. **Gradient Accumulation**:
111
+ Iterating in reverse topological order, `node->backward(grad_output)` computes intermediate derivatives and accumulates them into each input tensor's `grad` attribute.
112
+
113
+ ---
114
+
115
+ ### 2. Activation Checkpointing Algorithm (Re-Computation)
116
+
117
+ In deep Transformer models, storing all intermediate activations in VRAM causes Out-Of-Memory (OOM) failures.
118
+
119
+ > [!TIP]
120
+ > **Activation Checkpointing Mechanism**:
121
+ > Instead of keeping all intermediate activation tensors in VRAM during the forward pass, LiteTorch retains only the block input tensors. During the backward pass, LiteTorch automatically re-evaluates the block forward pass on-the-fly to re-compute activation tensors right when gradients are evaluated.
122
+
123
+ ```
124
+ [Standard Forward Pass]
125
+ Input ---> [Layer 1] ---> Act 1 ---> [Layer 2] ---> Act 2 ---> Loss
126
+ (All Act 1 & Act 2 must remain pinned in VRAM)
127
+
128
+ [Activation Checkpointing Pass]
129
+ Forward: Input ---> [Layer 1 & 2 under NoGradGuard] ---> Loss (Act 1 & Act 2 released from VRAM)
130
+ Backward: Input ---> [Re-compute Layer 1 & 2] ---> Evaluate Act 1 & 2 locally ---> Propagate Gradients
131
+ ```
132
+
133
+ ---
134
+
135
+ ### 3. FSDP & ZeRO-3 Distributed Parallelism Algorithm
136
+
137
+ LiteTorch implements **ZeRO-3 (Zero Redundancy Optimizer Stage 3)** and **Fully Sharded Data Parallel (FSDP)** to partition model states across $N$ GPUs.
138
+
139
+ #### Sharded State Categories:
140
+ - **Optimizer State Sharding**: Optimizer memory ($m, v$ in Adam) is sharded $\frac{1}{N}$ across GPUs.
141
+ - **Gradient Sharding**: Gradients are reduced via `Reduce-Scatter` and stored $\frac{1}{N}$ on respective owner GPUs.
142
+ - **Parameter Sharding**: Model parameters are sharded $\frac{1}{N}$ across GPUs.
143
+
144
+ ```mermaid
145
+ sequenceDiagram
146
+ participant GPU0 as GPU 0 (Owns Shard 0)
147
+ participant GPU1 as GPU 1 (Owns Shard 1)
148
+ Note over GPU0,GPU1: 1. Before Forward Pass
149
+ GPU0->>GPU1: All-Gather (Reconstruct full parameters for current layer)
150
+ Note over GPU0,GPU1: 2. Execute Forward & Release non-owned Parameter Shards
151
+ Note over GPU0,GPU1: 3. Before Backward Pass
152
+ GPU0->>GPU1: All-Gather (Reconstruct full parameters for gradient evaluation)
153
+ Note over GPU0,GPU1: 4. After Backward Pass
154
+ GPU0->>GPU1: Reduce-Scatter (Aggregate and shard gradients back to owner GPUs)
155
+ Note over GPU0,GPU1: 5. Local Optimizer Step on Sharded Parameters
156
+ ```
157
+
158
+ ---
159
+
160
+ ## System Console Commands
161
+
162
+ After installation, run benchmarks directly anywhere in your terminal without `./` or `python3` prefixes:
163
+
164
+ ```bash
165
+ # Run spiral dataset classification benchmark
166
+ demo_run.py
167
+
168
+ # Run Python bindings test suite
169
+ test_litetorch.py
170
+ ```
171
+
172
+ ### Auto-Detection vs Forced Fallback
173
+
174
+ ```bash
175
+ # Auto-detection (Prefers CUDA/ROCm -> OpenCL -> CPU):
176
+ demo_run.py
177
+
178
+ # Force OpenCL / CPU Testing Mode (For local testing without CUDA GPU):
179
+ LITETORCH_NO_NATIVE_GPU=1 demo_run.py
180
+ ```
181
+
182
+ ---
183
+
184
+ ## Beginner Code Examples
185
+
186
+ ### Example 1: Hardware Auto-Detection & Autograd (Python)
187
+
188
+ ```python
189
+ import litetorch as lt
190
+
191
+ device = lt.auto_device()
192
+ print("Selected Device:", device)
193
+
194
+ if lt.cuda.is_available():
195
+ print("Running on Native NVIDIA CUDA / AMD ROCm GPU!")
196
+ elif lt.is_gpu_available():
197
+ print("Running on OpenCL GPU!")
198
+ else:
199
+ print("Running on CPU!")
200
+
201
+ x = lt.Tensor.from_vector([1.0, 2.0, 3.0, 4.0], [2, 2], device, True)
202
+ y = lt.Tensor.from_vector([2.0, 0.5, 1.0, 2.0], [2, 2], device, True)
203
+
204
+ z = lt.Ops.add(x, y)
205
+ loss = lt.Ops.sum(z)
206
+
207
+ loss.backward()
208
+
209
+ print("Loss Value:", loss.item())
210
+ print("Gradient of Tensor x:", x.grad.to_vector())
211
+ ```
212
+
213
+ ---
214
+
215
+ ### Example 2: Neural Network Training Loop (Python)
216
+
217
+ ```python
218
+ import litetorch as lt
219
+
220
+ device = lt.auto_device()
221
+
222
+ x_data = lt.Tensor.from_vector([0.5, 1.5, 2.0, 3.0], [2, 2], device, False)
223
+ y_data = lt.Tensor.from_vector([1.0, 0.0], [2], device, False)
224
+
225
+ class NeuralNetwork(lt.nn.Module):
226
+ def __init__(self):
227
+ super().__init__()
228
+ self.fc1 = lt.nn.Linear(2, 8, True)
229
+ self.fc2 = lt.nn.Linear(8, 2, True)
230
+
231
+ def forward(self, x):
232
+ h = self.fc1.forward(x)
233
+ act = lt.Ops.relu(h)
234
+ return self.fc2.forward(act)
235
+
236
+ def parameters(self):
237
+ return self.fc1.parameters() + self.fc2.parameters()
238
+
239
+ model = NeuralNetwork()
240
+ optimizer = lt.optim.AdamW(model.parameters(), lr=0.01)
241
+
242
+ for epoch in range(1, 101):
243
+ optimizer.zero_grad()
244
+ out = model.forward(x_data)
245
+ loss = lt.Ops.cross_entropy_loss(out, y_data)
246
+ loss.backward()
247
+ optimizer.step()
248
+
249
+ if epoch % 20 == 0:
250
+ print(f"Epoch {epoch:3d} | Loss: {loss.item():.6f}")
251
+ ```
252
+
253
+ ---
254
+
255
+ ### Example 3: Memory-Optimized Activation Checkpointing (Python)
256
+
257
+ ```python
258
+ import litetorch as lt
259
+
260
+ device = lt.auto_device()
261
+
262
+ x = lt.Tensor.from_vector([1.0, 2.0, 3.0, 4.0], [4], device, True)
263
+
264
+ def heavy_layer(inp):
265
+ h = lt.Ops.mul(inp, inp)
266
+ return lt.Ops.add(h, inp)
267
+
268
+ output = lt.checkpoint(heavy_layer, x)
269
+ loss = lt.Ops.sum(output)
270
+
271
+ loss.backward()
272
+
273
+ print("Checkpointed Gradient:", x.grad.to_vector())
274
+ ```
275
+
276
+ ---
277
+
278
+ ## Benchmark Results
279
+
280
+ Training 300 epochs on 600-sample 3-class spiral dataset (`demo_run.py`):
281
+
282
+ | Metric | Measured Result | Technical Details |
283
+ | :--- | :--- | :--- |
284
+ | **Final Accuracy** | **100.00%** | Converged perfectly at Epoch 200 |
285
+ | **Final Loss** | **0.000804** | Loss dropped close to zero |
286
+ | **RAM Consumption** | **34.45 MB** | Extremely lightweight RSS RAM footprint |
287
+ | **Peak RAM** | **33.98 MB** | Maximum RAM usage throughout training |
288
+ | **Total CPU Time** | **12.50 seconds** | CPU execution time |
289
+ | **Wall-Clock Time** | **9.48 seconds** | Total end-to-end elapsed time |
290
+ | **Build Speed (`make -j8`)** | **< 0.3 seconds** | **480x faster** than legacy sequential compilation |
291
+
292
+ ---
293
+
294
+ ## License
295
+
296
+ LiteTorch is open-sourced under the **MIT License**.
297
+ "# Lt"
@@ -0,0 +1,275 @@
1
+ > [!WARNING]
2
+ > **THIS PROJECT IS NOT YET COMPLETE AND MAY CONTAIN SOME ERRORS OR IMPROVEMENTS. WE ARE WORKING TO FIX THEM. YOU CAN ALSO CONTRIBUTE TO THE FRAMEWORK!.**
3
+
4
+
5
+ # LiteTorch framework
6
+
7
+ LiteTorch is a lightweight, high-performance deep learning framework built natively in C++14 with seamless Python bindings via `pybind11`. Designed with an intuitive PyTorch-like API, LiteTorch features dynamic autograd graph execution, memory optimization via Activation Checkpointing, distributed training primitives (FSDP, ZeRO-3), and multi-backend hardware acceleration (NVIDIA CUDA, AMD ROCm, OpenCL, and multi-threaded CPU).
8
+
9
+ ---
10
+
11
+ ## Key Features
12
+
13
+ - **PyTorch-Style Hardware Auto-Detection**:
14
+ - Automatically senses and initializes **NVIDIA CUDA** (`nvcc` + cuBLAS/cuDNN) or **AMD ROCm/HIP** (`hipcc` + rocBLAS/MIOpen) when native GPUs are present.
15
+ - Seamlessly falls back to **OpenCL** or multi-threaded **CPU** execution on systems without native GPU drivers.
16
+ - **Dual API (C++ Core & Python Bindings)**:
17
+ - High-level Python interface: `import litetorch as lt`.
18
+ - Zero performance overhead with native C++14 execution underneath.
19
+ - **Dynamic Autograd Engine**:
20
+ - Reverse-mode automatic differentiation over Directed Acyclic Graphs (DAG).
21
+ - Topological Sort DAG traversal algorithm for precise gradient accumulation.
22
+ - **Advanced Memory Management**:
23
+ - **Activation Checkpointing**: Re-computes activations during backward passes to dramatically reduce VRAM footprint.
24
+ - **LRU Storage Eviction & Caching Allocator**: Smart memory pooling and automatic LRU swap between RAM and VRAM.
25
+ - **Distributed Training Primitives**:
26
+ - **Fully Sharded Data Parallel (FSDP)** & **ZeRO-3 Optimizer**: Shards parameters, gradients, and optimizer states across GPUs.
27
+ - Inter-node communication via NCCL (NVIDIA), RCCL (AMD), Shared Memory IPC (SHM), and TCP Socket fallback.
28
+ - **Fast Build System**: Multi-core parallel Makefile (`make -j$(nproc)`) with shared library caching (`liblitetorch.so`) enabling sub-second test runs.
29
+ - **System Console Commands**: Global command registration for executing `demo_run.py` or `test_litetorch.py` directly without `./` or `python3` prefixes.
30
+
31
+ ---
32
+
33
+ ## Quick Installation & Setup
34
+
35
+ ### 1. Auto-Install C++ Build Dependencies (Linux Auto-Installer)
36
+
37
+ Automated installer script for C++ dependencies on Linux (Ubuntu, Debian, RHEL, Fedora, Arch Linux):
38
+
39
+ ```bash
40
+ ./install_deps.sh
41
+ ```
42
+
43
+ *(For detailed OS-specific C++ & GPU toolkit installation guides, see [`REQUIREMENTS_CPP.md`](file:///home/notmerblx/Pictures/Litetorch/REQUIREMENTS_CPP.md))*
44
+
45
+ ### 2. Install Python Dependencies & Package
46
+
47
+ ```bash
48
+ python3 -m pip install -r requirements.txt
49
+ python3 -m pip install -e .
50
+ ```
51
+
52
+ After installation, you can `import litetorch as lt` or run `demo_run.py` directly anywhere in your shell!
53
+
54
+ ---
55
+
56
+ ## Architecture & Core Algorithms Breakdown
57
+
58
+ ### Layered System Architecture
59
+
60
+ ```mermaid
61
+ graph TD
62
+ A["Python Layer (import litetorch as lt)"] --> B["C++ Binding Layer (pybind11)"]
63
+ B --> C["LiteTorch High-Level API (Tensor, Ops, nn::Module, optim)"]
64
+ C --> D["Autograd & Memory Engine (DAG, Checkpointing, Caching Allocator)"]
65
+ D --> E["Distributed Engine (ProcessGroup, FSDP, ZeRO-3, NCCL/RCCL)"]
66
+ E --> F1["Backend 1: Native GPU (CUDA / ROCm cuBLAS/rocBLAS)"]
67
+ E --> F2["Backend 2: OpenCL Backend"]
68
+ E --> F3["Backend 3: Multi-Threaded CPU Engine"]
69
+ ```
70
+
71
+ ---
72
+
73
+ ### 1. Dynamic Autograd Graph & Topological Sort Algorithm
74
+
75
+ Tensor operations dynamically construct a **Directed Acyclic Graph (DAG)** where each `Tensor` acts as a Node holding a weak pointer (`std::weak_ptr<Node> creator`) to the operation that produced it.
76
+
77
+ ```mermaid
78
+ graph LR
79
+ X["Tensor X (Input)"] -->|mul| H1["Tensor H1"]
80
+ X -->|mul| H1
81
+ H1 -->|add| H2["Tensor H2 (Output)"]
82
+ X -->|add| H2
83
+ ```
84
+
85
+ #### Reverse-Mode Automatic Differentiation Workflow:
86
+ 1. **Topological Sort Traversal**:
87
+ When `loss->backward()` is called, the autograd engine executes a DFS or Kahn's algorithm to sort nodes from Output (Loss) back to Inputs.
88
+ 2. **Gradient Accumulation**:
89
+ Iterating in reverse topological order, `node->backward(grad_output)` computes intermediate derivatives and accumulates them into each input tensor's `grad` attribute.
90
+
91
+ ---
92
+
93
+ ### 2. Activation Checkpointing Algorithm (Re-Computation)
94
+
95
+ In deep Transformer models, storing all intermediate activations in VRAM causes Out-Of-Memory (OOM) failures.
96
+
97
+ > [!TIP]
98
+ > **Activation Checkpointing Mechanism**:
99
+ > Instead of keeping all intermediate activation tensors in VRAM during the forward pass, LiteTorch retains only the block input tensors. During the backward pass, LiteTorch automatically re-evaluates the block forward pass on-the-fly to re-compute activation tensors right when gradients are evaluated.
100
+
101
+ ```
102
+ [Standard Forward Pass]
103
+ Input ---> [Layer 1] ---> Act 1 ---> [Layer 2] ---> Act 2 ---> Loss
104
+ (All Act 1 & Act 2 must remain pinned in VRAM)
105
+
106
+ [Activation Checkpointing Pass]
107
+ Forward: Input ---> [Layer 1 & 2 under NoGradGuard] ---> Loss (Act 1 & Act 2 released from VRAM)
108
+ Backward: Input ---> [Re-compute Layer 1 & 2] ---> Evaluate Act 1 & 2 locally ---> Propagate Gradients
109
+ ```
110
+
111
+ ---
112
+
113
+ ### 3. FSDP & ZeRO-3 Distributed Parallelism Algorithm
114
+
115
+ LiteTorch implements **ZeRO-3 (Zero Redundancy Optimizer Stage 3)** and **Fully Sharded Data Parallel (FSDP)** to partition model states across $N$ GPUs.
116
+
117
+ #### Sharded State Categories:
118
+ - **Optimizer State Sharding**: Optimizer memory ($m, v$ in Adam) is sharded $\frac{1}{N}$ across GPUs.
119
+ - **Gradient Sharding**: Gradients are reduced via `Reduce-Scatter` and stored $\frac{1}{N}$ on respective owner GPUs.
120
+ - **Parameter Sharding**: Model parameters are sharded $\frac{1}{N}$ across GPUs.
121
+
122
+ ```mermaid
123
+ sequenceDiagram
124
+ participant GPU0 as GPU 0 (Owns Shard 0)
125
+ participant GPU1 as GPU 1 (Owns Shard 1)
126
+ Note over GPU0,GPU1: 1. Before Forward Pass
127
+ GPU0->>GPU1: All-Gather (Reconstruct full parameters for current layer)
128
+ Note over GPU0,GPU1: 2. Execute Forward & Release non-owned Parameter Shards
129
+ Note over GPU0,GPU1: 3. Before Backward Pass
130
+ GPU0->>GPU1: All-Gather (Reconstruct full parameters for gradient evaluation)
131
+ Note over GPU0,GPU1: 4. After Backward Pass
132
+ GPU0->>GPU1: Reduce-Scatter (Aggregate and shard gradients back to owner GPUs)
133
+ Note over GPU0,GPU1: 5. Local Optimizer Step on Sharded Parameters
134
+ ```
135
+
136
+ ---
137
+
138
+ ## System Console Commands
139
+
140
+ After installation, run benchmarks directly anywhere in your terminal without `./` or `python3` prefixes:
141
+
142
+ ```bash
143
+ # Run spiral dataset classification benchmark
144
+ demo_run.py
145
+
146
+ # Run Python bindings test suite
147
+ test_litetorch.py
148
+ ```
149
+
150
+ ### Auto-Detection vs Forced Fallback
151
+
152
+ ```bash
153
+ # Auto-detection (Prefers CUDA/ROCm -> OpenCL -> CPU):
154
+ demo_run.py
155
+
156
+ # Force OpenCL / CPU Testing Mode (For local testing without CUDA GPU):
157
+ LITETORCH_NO_NATIVE_GPU=1 demo_run.py
158
+ ```
159
+
160
+ ---
161
+
162
+ ## Beginner Code Examples
163
+
164
+ ### Example 1: Hardware Auto-Detection & Autograd (Python)
165
+
166
+ ```python
167
+ import litetorch as lt
168
+
169
+ device = lt.auto_device()
170
+ print("Selected Device:", device)
171
+
172
+ if lt.cuda.is_available():
173
+ print("Running on Native NVIDIA CUDA / AMD ROCm GPU!")
174
+ elif lt.is_gpu_available():
175
+ print("Running on OpenCL GPU!")
176
+ else:
177
+ print("Running on CPU!")
178
+
179
+ x = lt.Tensor.from_vector([1.0, 2.0, 3.0, 4.0], [2, 2], device, True)
180
+ y = lt.Tensor.from_vector([2.0, 0.5, 1.0, 2.0], [2, 2], device, True)
181
+
182
+ z = lt.Ops.add(x, y)
183
+ loss = lt.Ops.sum(z)
184
+
185
+ loss.backward()
186
+
187
+ print("Loss Value:", loss.item())
188
+ print("Gradient of Tensor x:", x.grad.to_vector())
189
+ ```
190
+
191
+ ---
192
+
193
+ ### Example 2: Neural Network Training Loop (Python)
194
+
195
+ ```python
196
+ import litetorch as lt
197
+
198
+ device = lt.auto_device()
199
+
200
+ x_data = lt.Tensor.from_vector([0.5, 1.5, 2.0, 3.0], [2, 2], device, False)
201
+ y_data = lt.Tensor.from_vector([1.0, 0.0], [2], device, False)
202
+
203
+ class NeuralNetwork(lt.nn.Module):
204
+ def __init__(self):
205
+ super().__init__()
206
+ self.fc1 = lt.nn.Linear(2, 8, True)
207
+ self.fc2 = lt.nn.Linear(8, 2, True)
208
+
209
+ def forward(self, x):
210
+ h = self.fc1.forward(x)
211
+ act = lt.Ops.relu(h)
212
+ return self.fc2.forward(act)
213
+
214
+ def parameters(self):
215
+ return self.fc1.parameters() + self.fc2.parameters()
216
+
217
+ model = NeuralNetwork()
218
+ optimizer = lt.optim.AdamW(model.parameters(), lr=0.01)
219
+
220
+ for epoch in range(1, 101):
221
+ optimizer.zero_grad()
222
+ out = model.forward(x_data)
223
+ loss = lt.Ops.cross_entropy_loss(out, y_data)
224
+ loss.backward()
225
+ optimizer.step()
226
+
227
+ if epoch % 20 == 0:
228
+ print(f"Epoch {epoch:3d} | Loss: {loss.item():.6f}")
229
+ ```
230
+
231
+ ---
232
+
233
+ ### Example 3: Memory-Optimized Activation Checkpointing (Python)
234
+
235
+ ```python
236
+ import litetorch as lt
237
+
238
+ device = lt.auto_device()
239
+
240
+ x = lt.Tensor.from_vector([1.0, 2.0, 3.0, 4.0], [4], device, True)
241
+
242
+ def heavy_layer(inp):
243
+ h = lt.Ops.mul(inp, inp)
244
+ return lt.Ops.add(h, inp)
245
+
246
+ output = lt.checkpoint(heavy_layer, x)
247
+ loss = lt.Ops.sum(output)
248
+
249
+ loss.backward()
250
+
251
+ print("Checkpointed Gradient:", x.grad.to_vector())
252
+ ```
253
+
254
+ ---
255
+
256
+ ## Benchmark Results
257
+
258
+ Training 300 epochs on 600-sample 3-class spiral dataset (`demo_run.py`):
259
+
260
+ | Metric | Measured Result | Technical Details |
261
+ | :--- | :--- | :--- |
262
+ | **Final Accuracy** | **100.00%** | Converged perfectly at Epoch 200 |
263
+ | **Final Loss** | **0.000804** | Loss dropped close to zero |
264
+ | **RAM Consumption** | **34.45 MB** | Extremely lightweight RSS RAM footprint |
265
+ | **Peak RAM** | **33.98 MB** | Maximum RAM usage throughout training |
266
+ | **Total CPU Time** | **12.50 seconds** | CPU execution time |
267
+ | **Wall-Clock Time** | **9.48 seconds** | Total end-to-end elapsed time |
268
+ | **Build Speed (`make -j8`)** | **< 0.3 seconds** | **480x faster** than legacy sequential compilation |
269
+
270
+ ---
271
+
272
+ ## License
273
+
274
+ LiteTorch is open-sourced under the **MIT License**.
275
+ "# Lt"
@@ -0,0 +1,36 @@
1
+ #ifndef LITETORCH_ALLOCATOR_H
2
+ #define LITETORCH_ALLOCATOR_H
3
+
4
+ #include <cstddef>
5
+ #include <mutex>
6
+ #include <map>
7
+
8
+ namespace litetorch {
9
+
10
+ class CachingAllocator {
11
+ public:
12
+ static CachingAllocator& get();
13
+
14
+ void* allocate_cpu(size_t size);
15
+ void free_cpu(void* ptr);
16
+
17
+ void* allocate_gpu(size_t size);
18
+ void free_gpu(void* ptr);
19
+
20
+ void empty_cache();
21
+
22
+ private:
23
+ CachingAllocator() = default;
24
+ ~CachingAllocator();
25
+
26
+ std::mutex mutex_;
27
+ std::multimap<size_t, void*> free_cpu_blocks_;
28
+ std::map<void*, size_t> allocated_cpu_blocks_;
29
+
30
+ std::multimap<size_t, void*> free_gpu_blocks_;
31
+ std::map<void*, size_t> allocated_gpu_blocks_;
32
+ };
33
+
34
+ }
35
+
36
+ #endif