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.
- litetorch-0.1.0/MANIFEST.in +8 -0
- litetorch-0.1.0/Makefile +32 -0
- litetorch-0.1.0/PKG-INFO +297 -0
- litetorch-0.1.0/README.md +275 -0
- litetorch-0.1.0/include/litetorch/allocator.h +36 -0
- litetorch-0.1.0/include/litetorch/amp.h +116 -0
- litetorch-0.1.0/include/litetorch/autograd.h +87 -0
- litetorch-0.1.0/include/litetorch/backend.h +84 -0
- litetorch-0.1.0/include/litetorch/checkpoint.h +18 -0
- litetorch-0.1.0/include/litetorch/cl_backend.h +227 -0
- litetorch-0.1.0/include/litetorch/context_parallel.h +15 -0
- litetorch-0.1.0/include/litetorch/continuous_batching.h +47 -0
- litetorch-0.1.0/include/litetorch/custom_ops.h +50 -0
- litetorch-0.1.0/include/litetorch/data.h +73 -0
- litetorch-0.1.0/include/litetorch/device.h +43 -0
- litetorch-0.1.0/include/litetorch/device_mesh.h +27 -0
- litetorch-0.1.0/include/litetorch/distributed.h +342 -0
- litetorch-0.1.0/include/litetorch/dtensor.h +41 -0
- litetorch-0.1.0/include/litetorch/fsdp.h +135 -0
- litetorch-0.1.0/include/litetorch/fsdp_wrapper.h +18 -0
- litetorch-0.1.0/include/litetorch/grad_scaler.h +33 -0
- litetorch-0.1.0/include/litetorch/guided_decoding.h +24 -0
- litetorch-0.1.0/include/litetorch/jit.h +182 -0
- litetorch-0.1.0/include/litetorch/llm_serving.h +62 -0
- litetorch-0.1.0/include/litetorch/memory_manager.h +51 -0
- litetorch-0.1.0/include/litetorch/nn.h +357 -0
- litetorch-0.1.0/include/litetorch/ops.h +113 -0
- litetorch-0.1.0/include/litetorch/optim.h +152 -0
- litetorch-0.1.0/include/litetorch/platform.h +38 -0
- litetorch-0.1.0/include/litetorch/quantization.h +146 -0
- litetorch-0.1.0/include/litetorch/serialization.h +23 -0
- litetorch-0.1.0/include/litetorch/tensor.h +158 -0
- litetorch-0.1.0/include/litetorch/thread_pool.h +161 -0
- litetorch-0.1.0/include/litetorch/zero3_optimizer.h +33 -0
- litetorch-0.1.0/litetorch.egg-info/PKG-INFO +297 -0
- litetorch-0.1.0/litetorch.egg-info/SOURCES.txt +168 -0
- litetorch-0.1.0/litetorch.egg-info/dependency_links.txt +1 -0
- litetorch-0.1.0/litetorch.egg-info/entry_points.txt +3 -0
- litetorch-0.1.0/litetorch.egg-info/not-zip-safe +1 -0
- litetorch-0.1.0/litetorch.egg-info/top_level.txt +1 -0
- litetorch-0.1.0/pyproject.toml +3 -0
- litetorch-0.1.0/setup.cfg +4 -0
- litetorch-0.1.0/setup.py +88 -0
- litetorch-0.1.0/src/autograd/autograd_engine.cpp +126 -0
- litetorch-0.1.0/src/autograd/checkpoint.cpp +66 -0
- litetorch-0.1.0/src/autograd/saved_tensor.cpp +37 -0
- litetorch-0.1.0/src/autograd.cpp +1 -0
- litetorch-0.1.0/src/backend/backend.cpp +316 -0
- litetorch-0.1.0/src/backend/gpu_native/common/gpu_common.h +182 -0
- litetorch-0.1.0/src/backend/gpu_native/elementwise/elementwise_ops.cu +502 -0
- litetorch-0.1.0/src/backend/gpu_native/entrypoint.cu +388 -0
- litetorch-0.1.0/src/backend/gpu_native/kernels.cu +1 -0
- litetorch-0.1.0/src/backend/gpu_native/kernels.hip +1 -0
- litetorch-0.1.0/src/backend/gpu_native/math/gemm.cu +173 -0
- litetorch-0.1.0/src/backend/gpu_native/math/reduction.cu +44 -0
- litetorch-0.1.0/src/backend/gpu_native/nn/flash_attention.cu +761 -0
- litetorch-0.1.0/src/backend/gpu_native/nn/nn_kernels.cu +1537 -0
- litetorch-0.1.0/src/backend/gpu_native/optim/optimizers.cu +124 -0
- litetorch-0.1.0/src/bindings/python_bindings.cpp +823 -0
- litetorch-0.1.0/src/cl_backend/cl_backend.cpp +655 -0
- litetorch-0.1.0/src/cl_backend/cl_functions.cpp +45 -0
- litetorch-0.1.0/src/cl_backend/cl_functions.h +56 -0
- litetorch-0.1.0/src/cl_backend/command_graph.cpp +48 -0
- litetorch-0.1.0/src/cl_backend.cpp +1 -0
- litetorch-0.1.0/src/data/dataloader.cpp +352 -0
- litetorch-0.1.0/src/data/dataset.cpp +41 -0
- litetorch-0.1.0/src/data.cpp +1 -0
- litetorch-0.1.0/src/distributed/context_parallel.cpp +79 -0
- litetorch-0.1.0/src/distributed/distributed.cpp +1287 -0
- litetorch-0.1.0/src/distributed/fsdp_wrapper.cpp +159 -0
- litetorch-0.1.0/src/distributed/shm_ipc.cpp +385 -0
- litetorch-0.1.0/src/llm/continuous_batching.cpp +200 -0
- litetorch-0.1.0/src/llm/guided_decoding.cpp +68 -0
- litetorch-0.1.0/src/llm/medusa.cpp +127 -0
- litetorch-0.1.0/src/memory_manager/allocator.cpp +99 -0
- litetorch-0.1.0/src/memory_manager/memory_manager.cpp +109 -0
- litetorch-0.1.0/src/memory_manager.cpp +1 -0
- litetorch-0.1.0/src/nn/activation.cpp +36 -0
- litetorch-0.1.0/src/nn/attention.cpp +95 -0
- litetorch-0.1.0/src/nn/container.cpp +79 -0
- litetorch-0.1.0/src/nn/convolution.cpp +73 -0
- litetorch-0.1.0/src/nn/dropout.cpp +47 -0
- litetorch-0.1.0/src/nn/embedding.cpp +160 -0
- litetorch-0.1.0/src/nn/linear.cpp +51 -0
- litetorch-0.1.0/src/nn/llm_serving.cpp +142 -0
- litetorch-0.1.0/src/nn/loss.cpp +41 -0
- litetorch-0.1.0/src/nn/moe.cpp +281 -0
- litetorch-0.1.0/src/nn/nn_utils.h +46 -0
- litetorch-0.1.0/src/nn/normalization.cpp +60 -0
- litetorch-0.1.0/src/nn/parallel_linear.cpp +119 -0
- litetorch-0.1.0/src/nn/pooling.cpp +29 -0
- litetorch-0.1.0/src/nn/qlora.cpp +70 -0
- litetorch-0.1.0/src/nn/transformer.cpp +72 -0
- litetorch-0.1.0/src/nn.cpp +1 -0
- litetorch-0.1.0/src/ops/activation.cpp +612 -0
- litetorch-0.1.0/src/ops/attention.cpp +474 -0
- litetorch-0.1.0/src/ops/cast.cpp +48 -0
- litetorch-0.1.0/src/ops/checkpoint.cpp +99 -0
- litetorch-0.1.0/src/ops/clip_grad.cpp +56 -0
- litetorch-0.1.0/src/ops/convolution.cpp +579 -0
- litetorch-0.1.0/src/ops/custom_ops.cpp +70 -0
- litetorch-0.1.0/src/ops/elementwise.cpp +1632 -0
- litetorch-0.1.0/src/ops/flash_decoding.cpp +118 -0
- litetorch-0.1.0/src/ops/fused_loss.cpp +198 -0
- litetorch-0.1.0/src/ops/jit.cpp +539 -0
- litetorch-0.1.0/src/ops/kernels.cl +2034 -0
- litetorch-0.1.0/src/ops/kernels.cpp +1 -0
- litetorch-0.1.0/src/ops/linear.cpp +379 -0
- litetorch-0.1.0/src/ops/loss.cpp +430 -0
- litetorch-0.1.0/src/ops/moe_ops.cpp +196 -0
- litetorch-0.1.0/src/ops/normalization.cpp +666 -0
- litetorch-0.1.0/src/ops/pooling.cpp +470 -0
- litetorch-0.1.0/src/ops/quantization_ops.cpp +106 -0
- litetorch-0.1.0/src/ops/ring_attention.cpp +160 -0
- litetorch-0.1.0/src/ops/rope.cpp +155 -0
- litetorch-0.1.0/src/ops/scaled_matmul.cpp +46 -0
- litetorch-0.1.0/src/ops/w8a8_ops.cpp +110 -0
- litetorch-0.1.0/src/ops.cpp +1 -0
- litetorch-0.1.0/src/optim/adam.cpp +73 -0
- litetorch-0.1.0/src/optim/adamw.cpp +104 -0
- litetorch-0.1.0/src/optim/adamw_8bit.cpp +107 -0
- litetorch-0.1.0/src/optim/adamw_fp8.cpp +129 -0
- litetorch-0.1.0/src/optim/cosine_annealing.cpp +24 -0
- litetorch-0.1.0/src/optim/grad_scaler.cpp +83 -0
- litetorch-0.1.0/src/optim/optim_utils.h +29 -0
- litetorch-0.1.0/src/optim/optimizer.cpp +15 -0
- litetorch-0.1.0/src/optim/rmsprop.cpp +61 -0
- litetorch-0.1.0/src/optim/sgd.cpp +81 -0
- litetorch-0.1.0/src/optim/steplr.cpp +18 -0
- litetorch-0.1.0/src/optim/zero3_optimizer.cpp +115 -0
- litetorch-0.1.0/src/optim.cpp +1 -0
- litetorch-0.1.0/src/serialization/optimizer.cpp +395 -0
- litetorch-0.1.0/src/serialization/parameters.cpp +167 -0
- litetorch-0.1.0/src/serialization.cpp +1 -0
- litetorch-0.1.0/src/tensor/storage.cpp +214 -0
- litetorch-0.1.0/src/tensor/tensor_core.cpp +1079 -0
- litetorch-0.1.0/src/tensor.cpp +1 -0
- litetorch-0.1.0/tests/__pycache__/demo_run.cpython-312.pyc +0 -0
- litetorch-0.1.0/tests/__pycache__/test_litetorch.cpython-312.pyc +0 -0
- litetorch-0.1.0/tests/__pycache__/test_litetorch.cpython-314.pyc +0 -0
- litetorch-0.1.0/tests/advanced_features_test.cpp +216 -0
- litetorch-0.1.0/tests/advanced_optimization_stage2_test.cpp +112 -0
- litetorch-0.1.0/tests/advanced_serving_test.cpp +159 -0
- litetorch-0.1.0/tests/allocator_test.cpp +37 -0
- litetorch-0.1.0/tests/attention_test.cpp +77 -0
- litetorch-0.1.0/tests/audit_bugs_test.cpp +86 -0
- litetorch-0.1.0/tests/audit_optimizations_test.cpp +35 -0
- litetorch-0.1.0/tests/backend_detection_test.cpp +70 -0
- litetorch-0.1.0/tests/checkpoint_test.cpp +32 -0
- litetorch-0.1.0/tests/demo_run.cpp +245 -0
- litetorch-0.1.0/tests/demo_run.py +156 -0
- litetorch-0.1.0/tests/distributed_test.cpp +96 -0
- litetorch-0.1.0/tests/double_backward_test.cpp +35 -0
- litetorch-0.1.0/tests/fusion_benchmark.cpp +64 -0
- litetorch-0.1.0/tests/gpt_training_test.cpp +123 -0
- litetorch-0.1.0/tests/gpu_allocator_test.cpp +32 -0
- litetorch-0.1.0/tests/gpu_cast_test.cpp +52 -0
- litetorch-0.1.0/tests/gpu_half_perf_test.cpp +200 -0
- litetorch-0.1.0/tests/jit_test.cpp +102 -0
- litetorch-0.1.0/tests/llm_optimization_test.cpp +187 -0
- litetorch-0.1.0/tests/memory_leak_test.cpp +85 -0
- litetorch-0.1.0/tests/new_features_test.cpp +546 -0
- litetorch-0.1.0/tests/production_features_test.cpp +306 -0
- litetorch-0.1.0/tests/production_optimization_test.cpp +166 -0
- litetorch-0.1.0/tests/scale_upgrade_test.cpp +105 -0
- litetorch-0.1.0/tests/stage4_test.cpp +125 -0
- litetorch-0.1.0/tests/stress_test.cpp +324 -0
- litetorch-0.1.0/tests/test_litetorch.py +59 -0
- litetorch-0.1.0/tests/transformer_test.cpp +125 -0
litetorch-0.1.0/Makefile
ADDED
|
@@ -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)
|
litetorch-0.1.0/PKG-INFO
ADDED
|
@@ -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
|