From 6c88e3d13808d738bead791529b070b9658ec3c4 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 13 Aug 2026 07:21:54 -0600 Subject: [PATCH] Add layer offloading to omni captioner --- .../captioner/BaseCaptioner.py | 4 + .../captioner/Qwen3OmniCaptioner.py | 109 ++++++++++++++++-- ui/src/components/CaptionSimpleJob.tsx | 24 ++++ ui/src/helpers/captionOptions.ts | 7 +- ui/src/types.ts | 2 + version.py | 2 +- 6 files changed, 136 insertions(+), 12 deletions(-) diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py index edfebb4..0333623 100644 --- a/extensions_built_in/captioner/BaseCaptioner.py +++ b/extensions_built_in/captioner/BaseCaptioner.py @@ -50,6 +50,10 @@ class CaptionConfig: # safety; raise it to saturate a large GPU. self.batch_size = kwargs.get("batch_size", 1) self.num_workers = kwargs.get("num_workers", 3) + # stream weights from CPU per layer instead of keeping them resident + # (low-vram machines); percent is the fraction of linears offloaded + self.layer_offloading = kwargs.get("layer_offloading", False) + self.layer_offloading_percent = kwargs.get("layer_offloading_percent", 1.0) class BaseCaptioner(BaseExtensionProcess): diff --git a/extensions_built_in/captioner/Qwen3OmniCaptioner.py b/extensions_built_in/captioner/Qwen3OmniCaptioner.py index c502f81..868ddc8 100644 --- a/extensions_built_in/captioner/Qwen3OmniCaptioner.py +++ b/extensions_built_in/captioner/Qwen3OmniCaptioner.py @@ -204,21 +204,59 @@ class ConvRot8Experts(torch.nn.Module): persistent=False, ) + # device the streamed experts should land on when the banks themselves + # stay in system RAM (low-vram layer offloading); None = banks resident + offload_device = None + + def enable_offload(self, device): + """Keep the int8 banks in (pinned) system RAM; forward streams only + the routed experts' rows to the GPU per layer call.""" + self.offload_device = device + try: + self.gate_up_q = self.gate_up_q.pin_memory() + self.down_q = self.down_q.pin_memory() + self.gate_up_s = self.gate_up_s.pin_memory() + self.down_s = self.down_s.pin_memory() + except RuntimeError: + pass # pinning is a speed optimization only; pageable still works + # the hadamard matrices are tiny — keep them resident + self.gate_up_h = self.gate_up_h.to(device) + self.down_h = self.down_h.to(device) + @staticmethod def _rotate(w, h, rot): shape = w.shape return (w.reshape(-1, shape[-1] // rot, rot) @ h).reshape(shape) + def _gather(self, qdata, scales_u8, hit): + """Expert rows + scales for the hit indices, on the compute device.""" + if self.offload_device is not None and qdata.device.type == "cpu": + # each expert's rows are a contiguous view of the pinned bank, so + # slice-copies DMA straight to the GPU with zero CPU-side gather + # work (a CPU index_select here memcpy'd ~2GB/token on all cores) + hit_list = hit.tolist() if torch.is_tensor(hit) else list(hit) + scales = scales_u8.view(torch.float32) + q = torch.stack( + [qdata[i].to(self.offload_device, non_blocking=True) for i in hit_list] + ) + s = torch.stack( + [scales[i].to(self.offload_device, non_blocking=True) for i in hit_list] + ) + return q, s + return qdata[hit], scales_u8.view(torch.float32)[hit] + def _dequant(self, qdata, scales_u8, h, rot, i): # scales are [E, out, 1]; rotation is self-inverse along the in dim - scales = scales_u8.view(torch.float32) - w = qdata[i].float() * scales[i] + q, s = self._gather( + qdata, scales_u8, i.reshape(1) if torch.is_tensor(i) else torch.tensor([i]) + ) + w = q[0].float() * s[0] return self._rotate(w, h, rot).to(self.out_dtype) def _dequant_batch(self, qdata, scales_u8, h, rot, hit, dtype): """Dequantize the hit experts in one shot: [n_hit, out, in].""" - scales = scales_u8.view(torch.float32) - w = qdata[hit].float() * scales[hit] + q, s = self._gather(qdata, scales_u8, hit) + w = q.float() * s return self._rotate(w, h, rot).to(dtype) def forward(self, hidden_states, top_k_index, top_k_weights): @@ -255,7 +293,12 @@ class ConvRot8Experts(torch.nn.Module): del w_gate_up h = F.silu(gate) * up w_down = self._dequant_batch( - self.down_q, self.down_s, self.down_h, self.down_rot, flat, hidden_states.dtype + self.down_q, + self.down_s, + self.down_h, + self.down_rot, + flat, + hidden_states.dtype, ) out = torch.bmm(h, w_down.transpose(1, 2)).squeeze(1) del w_down @@ -318,11 +361,17 @@ class ConvRot8Experts(torch.nn.Module): top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) current_state = hidden_states[token_idx] w_gate_up = self._dequant( - self.gate_up_q, self.gate_up_s, self.gate_up_h, self.gate_up_rot, expert_idx + self.gate_up_q, + self.gate_up_s, + self.gate_up_h, + self.gate_up_rot, + expert_idx, ) gate, up = F.linear(current_state, w_gate_up).chunk(2, dim=-1) current_hidden_states = F.silu(gate) * up - w_down = self._dequant(self.down_q, self.down_s, self.down_h, self.down_rot, expert_idx) + w_down = self._dequant( + self.down_q, self.down_s, self.down_h, self.down_rot, expert_idx + ) current_hidden_states = F.linear(current_hidden_states, w_down) current_hidden_states = ( current_hidden_states * top_k_weights[token_idx, top_k_pos, None] @@ -482,6 +531,14 @@ class Qwen3OmniCaptioner(BaseCaptioner): model.generation_config.pad_token_id = 151643 model.generation_config.eos_token_id = [151645, 151643] + # built from config, so no sampling defaults were loaded; greedy decode + # falls into repetition loops on long captions (A-B-A-B forever on + # low-motion clips). Qwen's recommended sampling for the Qwen3 family: + model.generation_config.do_sample = True + model.generation_config.temperature = 0.7 + model.generation_config.top_p = 0.8 + model.generation_config.top_k = 20 + model.generation_config.repetition_penalty = 1.05 # swap the slow bf16 Conv3d patch_embed for an equivalent fast linear patch_qwen_vl_patch_embed(model) @@ -493,6 +550,33 @@ class Qwen3OmniCaptioner(BaseCaptioner): ) self.model = model + if self.caption_config.layer_offloading: + from toolkit.memory_management import MemoryManager + + self.print_and_status_update( + " - layer offloading enabled: expert banks stay in system RAM, " + "linears stream per layer" + ) + # expert banks: stay in system RAM, stream routed experts per call + for module in model.modules(): + if isinstance(module, ConvRot8Experts): + module.enable_offload(self.device_torch) + # everything the manager doesn't classify must ride to the GPU as + # unmanaged: the output head, the MoE routers (bare-parameter + # modules doing F.linear directly), and buffer-only modules + ignore = [model.lm_head] + ignore += [ + m + for m in model.modules() + if m.__class__.__name__ == "SinusoidsPositionEmbedding" + or m.__class__.__name__.endswith("TopKRouter") + ] + MemoryManager.attach( + model, + self.device_torch, + offload_percent=self.caption_config.layer_offloading_percent, + ignore_modules=ignore, + ) self.model.to(self.device_torch) self.processor = AutoProcessor.from_pretrained(BASE_REPO) flush() @@ -578,7 +662,11 @@ class Qwen3OmniCaptioner(BaseCaptioner): input_len = inputs["input_ids"].shape[1] gen_kwargs["max_length"] = STATIC_MAX_LENGTH gen_kwargs["stopping_criteria"] = StoppingCriteriaList( - [MaxLengthCriteria(max_length=input_len + self.caption_config.max_new_tokens)] + [ + MaxLengthCriteria( + max_length=input_len + self.caption_config.max_new_tokens + ) + ] ) else: gen_kwargs["max_new_tokens"] = self.caption_config.max_new_tokens @@ -690,6 +778,11 @@ class Qwen3OmniCaptioner(BaseCaptioner): batch sizes. First video per batch shape is slow (compile warmup).""" if not self.caption_config.compile: return + if self.caption_config.layer_offloading: + # cuda graphs need every tensor GPU-resident; offloaded weights + # live in system RAM, so the compiled decode path cannot capture + print("[AITK] layer offloading is on; skipping compiled decode.") + return import importlib.util if importlib.util.find_spec("triton") is None: diff --git a/ui/src/components/CaptionSimpleJob.tsx b/ui/src/components/CaptionSimpleJob.tsx index 000ffff..f7866df 100644 --- a/ui/src/components/CaptionSimpleJob.tsx +++ b/ui/src/components/CaptionSimpleJob.tsx @@ -4,6 +4,7 @@ import { CreatableSelectInput, FormGroup, SelectInput, + SliderInput, TextAreaInput, TextInput, } from '@/components/formInputs'; @@ -208,6 +209,29 @@ const CaptionSimpleJob: React.FC = ({ jobConfig, setJobConfig, gpuIDs, se onChange={value => setJobConfig(value, 'config.process[0].caption.thinking')} /> )} + {additionalSections.includes('caption.layer_offloading') && ( + <> + setJobConfig(value, 'config.process[0].caption.layer_offloading')} + /> + {jobConfig.config.process[0].caption.layer_offloading && ( +
+ + setJobConfig(value * 0.01, 'config.process[0].caption.layer_offloading_percent') + } + min={0} + max={100} + step={1} + /> +
+ )} + + )} diff --git a/ui/src/helpers/captionOptions.ts b/ui/src/helpers/captionOptions.ts index e3a0af2..a1c3962 100644 --- a/ui/src/helpers/captionOptions.ts +++ b/ui/src/helpers/captionOptions.ts @@ -1,7 +1,7 @@ import { GroupedSelectOption, SelectOption } from "@/types"; type CaptionGroup = 'image' | 'music' | 'video'; -type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.thinking' | 'caption.batch_size'; +type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading'; export interface CaptionOption { name: string; @@ -130,6 +130,7 @@ export const captionerTypes: CaptionOption[] = [ 'caption.max_res', 'caption.max_new_tokens', 'caption.batch_size', + 'caption.layer_offloading', ], }, { @@ -179,8 +180,8 @@ export const groupedCaptionerTypes: GroupedSelectOption[] = captionerTypes.reduc export const quantizationOptions: SelectOption[] = [ { value: '', label: '- NONE -' }, - { value: 'float8', label: 'float8 (default)' }, - { value: 'convrot8', label: '8bit convrot' }, + { value: 'float8', label: 'float8' }, + { value: 'convrot8', label: '8bit convrot (default)' }, { value: 'convrot4', label: '4bit convrot (nvfp4)' }, { value: 'convrotint7', label: '7bit convrot' }, { value: 'convrotint6', label: '6bit convrot' }, diff --git a/ui/src/types.ts b/ui/src/types.ts index eb28747..f1de3fa 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -299,6 +299,8 @@ export interface CaptionProcessConfig { caption_extension?: string; thinking?: boolean; batch_size?: number; + layer_offloading?: boolean; + layer_offloading_percent?: number; } } diff --git a/version.py b/version.py index b019ced..74a18e7 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.12.17" +VERSION = "0.12.18"