Skip to content

Commit b3fa345

Browse files
aki916faki916f
committed
add global-ptq contents
1 parent 600d2d2 commit b3fa345

21 files changed

Lines changed: 9262 additions & 0 deletions

‎global_ptq/README.md‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
## GlobalPTQ
2+
3+
`GlobalPTQ` globally optimises quantization parameters via KL-divergence distillation from a full-precision teacher model.
4+
It supports both continuous parameters (scales, zeros, scaling factors) and discrete parameters (GPTQ integer weights, DBF binary matrices) with Straight-Through Estimator (STE) gradients.
5+
6+
### Single-GPU
7+
8+
```python
9+
from onecomp import Runner, ModelConfig, GPTQ
10+
from onecomp_globalptq import GlobalPTQ
11+
12+
model_config = ModelConfig(model_id="meta-llama/Llama-2-7b-hf", device="cuda:0")
13+
gptq = GPTQ(wbits=4, groupsize=128)
14+
15+
global_ptq = GlobalPTQ(
16+
epochs=5,
17+
gptq_lr=1e-5,
18+
gptq_optimize_intweight=True, # discrete integer-weight optimisation
19+
use_sam=True, # Sharpness-Aware Minimisation
20+
num_calibration_samples=128,
21+
max_length=2048,
22+
)
23+
24+
runner = Runner(
25+
model_config=model_config,
26+
quantizer=gptq,
27+
post_processes=[global_ptq],
28+
)
29+
runner.run()
30+
```
31+
32+
### Multi-GPU with DeepSpeed (GlobalPTQDistributed)
33+
34+
For large models that do not fit on a single GPU:
35+
36+
```python
37+
from onecomp import Runner, ModelConfig, GPTQ
38+
from onecomp_globaptq import GlobalPTQDistributed
39+
40+
model_config = ModelConfig(model_id="meta-llama/Llama-2-7b-hf", device="cuda:0")
41+
gptq = GPTQ(wbits=4, groupsize=128)
42+
43+
global_ptq = GlobalPTQDistributed(
44+
epochs=5,
45+
gptq_lr=1e-5,
46+
gptq_optimize_intweight=True,
47+
deepspeed_config="ds_zero2.json",
48+
num_calibration_samples=128,
49+
max_length=2048,
50+
)
51+
52+
runner = Runner(
53+
model_config=model_config,
54+
quantizer=gptq,
55+
post_processes=[global_ptq],
56+
)
57+
runner.run()
58+
```
59+
60+
Launch with `torchrun`:
61+
62+
```bash
63+
torchrun --nproc_per_node=2 my_script.py
64+
```
65+
66+
### Key Parameters
67+
68+
#### Continuous Parameter Optimisation
69+
70+
| Parameter | Default | Description |
71+
|-----------|---------|-------------|
72+
| `epochs` | `5` | Number of distillation epochs |
73+
| `gptq_lr` | `1e-5` | Learning rate for GPTQ scales/zeros |
74+
| `dbf_lr` | `5e-5` | Learning rate for DBF scaling parameters |
75+
| `temperature` | `1.0` | Softmax temperature for KL divergence |
76+
| `num_calibration_samples` | `128` | Number of calibration samples |
77+
| `max_length` | `2048` | Maximum sequence length |
78+
| `use_gradient_checkpointing` | `True` | Reduce GPU memory via recomputation |
79+
| `early_stopping_patience` | `0` | Stop early if KL does not improve (0 = disabled) |
80+
| `use_mixed_precision` | `False` | Enable BF16 autocast |
81+
| `grad_accum_steps` | `1` | Gradient accumulation steps |
82+
83+
#### Discrete Parameter Optimisation
84+
85+
| Parameter | Default | Description |
86+
|-----------|---------|-------------|
87+
| `gptq_optimize_intweight` | `False` | Optimise GPTQ integer weights via STE |
88+
| `gptq_intweight_lr` | `1e-4` | Learning rate for integer-weight parameters |
89+
| `optimize_binary` | `False` | Optimise DBF binary matrices via sign-STE |
90+
| `ste_k` | `100.0` | Smoothness for GPTQ integer-weight STE rounding |
91+
92+
#### Advanced Optimisation Techniques
93+
94+
| Parameter | Default | Description |
95+
|-----------|---------|-------------|
96+
| `use_sam` | `False` | Sharpness-Aware Minimisation |
97+
| `sam_rho` | `0.02` | SAM perturbation radius |
98+
| `use_ema` | `False` | Exponential Moving Average of parameters |
99+
| `ema_decay` | `0.99` | EMA decay rate |
100+
| `use_lookahead` | `False` | Lookahead optimiser wrapper |
101+
| `lookahead_k` | `5` | Lookahead sync interval |
102+
| `lookahead_alpha` | `0.5` | Lookahead interpolation weight |
103+
| `use_fisher_lr` | `False` | Fisher-information-adaptive per-layer LR |
104+
| `fisher_n_samples` | `4` | Samples for Fisher diagonal estimation |
105+
| `use_entropy_reg` | `False` | Entropy regularisation on weight distributions |
106+
| `entropy_lambda` | `0.1` | Entropy regularisation strength |
107+
| `use_inter_loss` | `False` | Intermediate-layer cosine alignment loss |
108+
| `lambda_inter` | `10.0` | Intermediate loss weight |
109+
| `use_progressive_unfreeze` | `False` | Gradually unfreeze layers from output to input |
110+
111+
#### GlobalPTQDistributed Additional Parameters
112+
113+
| Parameter | Default | Description |
114+
|-----------|---------|-------------|
115+
| `deepspeed_config` | `None` | Path to DeepSpeed config JSON |
116+
| `w_distill` | `1.0` | Weight for KL distillation loss |
117+
| `w_ntp` | `0.0` | Weight for next-token prediction loss |
118+
| `bf16` | `True` | Enable bfloat16 training |
119+
| `per_device_train_batch_size` | `1` | Batch size per GPU |
120+
| `gradient_accumulation_steps` | `1` | Gradient accumulation steps |
121+
122+
## Development
123+
124+
```bash
125+
pip install -e ".[dev]"
126+
# or, if using uv
127+
uv pip install -e ".[dev]"
128+
```
129+
130+
## License
131+
132+
Proprietary. See LICENSE for details.

‎global_ptq/example/example_global_ptq.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
"""
2+
Example: Global PTQ with OneComp
3+
4+
Copyright 2025-2026 Fujitsu Ltd.
5+
6+
"""
7+
8+
from onecomp import Runner, ModelConfig, GPTQ, CalibrationConfig, setup_logger
9+
from onecomp_globalptq import GlobalPTQ
10+
11+
MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
12+
13+
setup_logger()
14+
model_config = ModelConfig(model_id=MODEL_ID, device="cuda:0")
15+
calibration_config = CalibrationConfig(max_length=512, num_calibration_samples=128)
16+
gptq = GPTQ(wbits=4, groupsize=128)
17+
18+
runner = Runner(
19+
model_config=model_config,
20+
calibration_config=calibration_config,
21+
quantizer=gptq,
22+
post_processes=[GlobalPTQ()],
23+
)
24+
runner.run()

‎global_ptq/onecomp_globalptq/__init__.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
"""
2+
OneComp Global PTQ.
3+
4+
Copyright 2025-2026 Fujitsu Ltd.
5+
6+
Author: Keiji Kimura
7+
8+
"""
9+
10+
from .__version__ import __version__
11+
from .global_ptq import GlobalPTQ, GlobalPTQDistributed
12+
13+
__all__ = [
14+
"GlobalPTQ",
15+
"GlobalPTQDistributed",
16+
]

‎global_ptq/onecomp_globalptq/__version__.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
"""
2+
3+
Copyright 2025-2026 Fujitsu Ltd.
4+
5+
Author: Keiji Kimura
6+
7+
"""
8+
9+
__version__ = "0.1.0"

‎global_ptq/onecomp_globalptq/global_ptq/__init__.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
"""
2+
Global PTQ — model-wide post-training quantization.
3+
4+
Provides the full-featured ``GlobalPTQ`` and ``GlobalPTQDistributed``
5+
including discrete parameter optimisation, SAM, EMA, Lookahead,
6+
Fisher-adaptive LR, and other advanced techniques.
7+
8+
Copyright 2025-2026 Fujitsu Ltd.
9+
10+
Authors: Yoshiyuki Ishii, Keiji Kimura, Yuma Ichikawa
11+
12+
"""
13+
14+
from .global_ptq import GlobalPTQ
15+
from .global_ptq_distributed import GlobalPTQDistributed
16+
17+
__all__ = ["GlobalPTQ", "GlobalPTQDistributed"]

‎global_ptq/onecomp_globalptq/global_ptq/_core/__init__.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
"""
2+
Internal implementation for Global PTQ.
3+
4+
Copyright 2025-2026 Fujitsu Ltd.
5+
6+
Authors: Yoshiyuki Ishii, Keiji Kimura, Yuma Ichikawa
7+
8+
"""

0 commit comments

Comments
 (0)