Geta is very slow

[sumeshp99]

Your work seems very impressive. That's why I tried it on yolov8 model. I took some idea from the sanity check part on how to integrate it with yolov8.

Issue I am facing is that training seems to be very slow, as compared to normal training.

I implemented my code like this

from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.engine.trainer import BaseTrainer
import torch
from only_train_once import OTO
from only_train_once.quantization.quant_model import model_to_quantize_model
from only_train_once.quantization.quant_layers import QuantizationMode

class GetaTrainer(DetectionTrainer):
    def setup_model(self):
        # Call the base logic to build and move model
        ckpt = super().setup_model()

        # Now perform model‐modification (quant + prune setup)
        self.model = model_to_quantize_model(
            self.model,
            quant_mode=QuantizationMode.WEIGHT_AND_ACTIVATION,
            num_bits=16
        )

        dummy_input = torch.randn(1, 3, self.args.imgsz, self.args.imgsz)
        self.oto = OTO(model=self.model, dummy_input=dummy_input)

        self.oto.mark_unprunable_by_param_names([
            'model.22.dfl.conv.weight',
            'model.22.cv3.2.2.weight',
            'model.22.cv2.2.2.weight',
            'model.22.cv2.1.2.weight',
            'model.22.cv3.1.2.weight',
            'model.22.cv3.0.2.weight',
            'model.22.cv2.0.2.weight'
        ])

        for node_group in self.oto._graph.node_groups.values():
            for node in node_group:
                if node.op_name == 'slice':
                    node_group.is_prunable = False

        # self.oto.visualize(view=False, out_dir='.')
        return ckpt

    def build_optimizer(self, model, name="auto", lr=None, momentum=None, decay=None, iterations=None):
        # Use GETA to build optimizer based on the modified model
        optimizer = self.oto.geta(
            variant="sgd",
            lr=1e-3,
            lr_quant=1e-3,
            first_momentum=0.9,
            weight_decay=1e-4,
            target_group_sparsity=0.5,
            start_projection_step=5 * len(self.train_loader),
            projection_periods=5,
            projection_steps=2 * len(self.train_loader),
            start_pruning_step=10 * len(self.train_loader),
            pruning_periods=5,
            pruning_steps=10 * len(self.train_loader),
            bit_reduction=2,
            min_bit_wt=4,
            max_bit_wt=16
        )
        return optimizer

    def save_model(self):
        self.oto.construct_subnet(out_dir="./")

args = dict(model="yolov8n.pt", data="VOC.yaml", epochs=50, amp=False, batch=32, workers=16)
a = GetaTrainer(overrides=args)
a.train()

Normal speed I am getting for this is 5.5it/s, but after using geta, speeds I am getting is around 0.3it/s. Also the memory usage is around 12G. I assumed that training speed would not be effected much since it was mentioned that this was used to train BERT on squad.

So I want to ask, am I doing anything wrong or is this the expected?

Additional context: I am doing this training on an A100 20G MiG.

[aponte411]

sumeshp99 This is expected with your configuration. A few stacked sources of overhead:

1. SYMMETRIC_NONLINEAR is the default but SYMMETRIC_LINEAR is faster and we've found it works better in practice. Pass quant_type=QuantizationType.SYMMETRIC_LINEAR to model_to_quantize_model. Looking at quant_layers.py, SymQuantizerNonLinear runs extra torch.exp(t * torch.log(...)) ops in both forward and backward that SymQuantizerLinear skips.

2. WEIGHT_AND_ACTIVATION roughly doubles per-layer overhead. Every activation also goes through the custom autograd.Function, and for YOLOv8 activations are spatial (B x C x H x W), much bigger than weights. If your goal is INT8 deployment, weight-only QAT is faster.

3. Optimizer overhead is small. The quant ops add real per-step compute that the optimizer analysis doesn't cover. GETA's optimizer adds |G|·P amortized over training (P = pruning periods) plus O(L) per-step work (number of quantized layers, ~60–100). Both are small vs the standard per-step backward, which scales with N (millions). HESSO paper Appendix C, Table 10 has the related analysis (~0.015% extra space on Phi2). Separately, GETA's quant forward/backward adds elementwise ops scaling with the total number of parameters that GETA converted to QuantizeLinear/QuantizeConv2D layers (N_W), which is per step, plus N_A in WEIGHT_AND_ACTIVATION mode. That extra is not in the optimizer accounting and is the real per-step cost. Points 1, 2 and 4 help reduce it.

• https://arxiv.org/abs/2409.09085

4. The custom autograd.Function is Python-side. Even SYMMETRIC_LINEAR runs in Python per layer per step, and is a graph break for torch.compile. No fused CUDA kernel exists. This is inherent to the framework as currently written and is the dominant per-step cost.
   Recommended starting config:

self.model = model_to_quantize_model(
    self.model,
    quant_type=QuantizationType.SYMMETRIC_LINEAR,
    quant_mode=QuantizationMode.WEIGHT_ONLY,
    num_bits=16,  # keep as initial bit width; GETA will reduce per layer
)