WIP - updated for qwen3.6 + vastai provisioning by johndpope · Pull Request #27 · pengzhangzhi/Open-dLLM

johndpope · 2026-05-19T04:29:53Z

https://wandb.ai/snoozie/open-dllm-27b/runs/qwen3.6-27b-repr-align-cloud?nw=nwusersnoozie

VRAM: 17.9 GB / 97.9 GB each (18%) — ZeRO-3 params partitioned, optimizer on CPU

GPU compute: 0% — in between ZeRO-3 all-gather bursts (polled at an idle moment)
Power: 8W / 14W — confirming idle between compute bursts (TDP is 300W each)
CPU: workers at 98-99% — actively doing the ZeRO-3 forward/backward pass
RAM: 18.3-18.5% of 1.1 TB ≈ 207 GB each — that's the CPU optimizer states

The 0% GPU utilization is a polling artifact — nvidia-smi samples instantaneously and ZeRO-3 GPU bursts are short relative to the PCIe transfer pauses between layers. If you sampled 100 times you'd see it spike
to 80-100% during each layer's matmul, then drop back to 0% while the next layer's params are gathered over PCIe.

Root issue: 17.9 GB used of 97.9 GB = 81% of VRAM wasted. The GPU is massively underutilised because the optimizer lives on CPU and params travel over PCIe (16 GB/s) for each of 64 layers. This is the ZeRO-3 +
CPU offload tax.

4× H100 on NVLink would fix this — full VRAM, no PCIe bottleneck, 900 GB/s all-gather.

loss: 9.29 | grad_norm: 5060.58 | lr: 5.00e-05
Step 1/200 in 8:39 | ETA ~28h 43m total

Loss 9.29 matches the previous successful run exactly — confirms the new repr_align_sub_sample_ratio: 0.25 code is working correctly
grad_norm 5060 is high but expected at step 1 (model just started adapting its bidirectional attention)
8m39s/step × 199 remaining = ~28.7 hours to complete all 200 steps

- Add qwen3_5 dense model with hybrid linear/full attention (Gated DeltaNet) - Add qwen3_5_moe MoE variant (256 experts, shared expert, expert parallelism) - Bump transformers>=4.57.6 for qwen3_next compat - Add flash-linear-attention>=0.2.0 for Gated DeltaNet kernels - Register models in __init__.py and auto-discovery via ModelClass - Add model configs for Qwen3.6-27B and Qwen3.6-35B-A3B - Add pretrain configs for both variants - Add is_fla_available() import utility - Update CLAUDE.md with new model families

- Add LDLMAutoencoder with Perceiver-based latent encoder/decoder - Add DiffusionHead (lightweight DiT) for latent-space diffusion - Add AdaptiveTimestepSampler from LDLM paper (Section 5.3) - Add LDLMTrainer class for training loop integration - Add LDLMArguments config dataclass - Register ldlm field in ModelArguments for YAML parsing - Add pretrain configs for Qwen3.6-27B and Qwen3.6-35B-A3B LDLM training - All syntax verified (ast.parse passes)

- Add tasks/train_ldlm.py: dedicated training script with FSDP, checkpointing, wandb logging, and LDLM-specific forward pass - Update LDLMTrainer with config-driven wandb logging: - Core losses (diffusion, reconstruction, total) - Latent statistics (norm, std, histograms) - Adaptive sampler diagnostics (loss min/max/range) - Generation evaluation (perplexity, entropy) - Fix transformers 5.8.0 API breaks: - is_safetensors_available -> direct try/except import - AutoModelForVision2Seq -> AutoModelForImageTextToText - no_init_weights removed (redundant with init_empty_weights) - Fix all references in loader.py, seed_omni, module_utils

- Add log_interval, gen_eval_interval, log_latent_histograms, log_samples to LDLMArguments dataclass - Add logging fields to both pretrain YAML configs (27B + 35B) - Make trainer_patch.py use config-driven histogram interval - Make train_ldlm.py use config-driven log intervals instead of hardcoded 50/500

Add configs/wandb_dashboard_ldlm.json with 8 pre-configured panels: - Core loss breakdown (diffusion vs reconstruction) - Latent space health (norm/std for collapse detection) - Sampler diagnostics (loss min/max/range) - Training signals (LR, grad norm, warmup) - Generation quality (PPL, entropy) - Latent histograms (mean, std distributions) - Decoder noise tracking

Add prune_checkpoints() to train_ldlm.py — runs after every Checkpointer.save() on rank 0, deletes all but the 2 most recent global_step_* directories to save disk space.

- MoE config has hidden_size/vocab_size inside text_config sub-object - Add robust accessor that checks both flat and nested configs - Fix autoencoder.py and train_ldlm.py to use self._vocab_size - Fix config YAMLs: remove 'type' field (not a recognized arg), remove expert_parallel_size=8 (only 2 GPUs), use local test dataset

- MoE model has hidden_size=2048, nhead=12 doesn't divide evenly - Auto-select nhead from [16,12,8,4,2] based on what divides dim - 2048/16=128 for Qwen3.6-35B-A3B, 5120/16=320 for 27B dense

- 35B MoE model OOMs when .cuda() moves all params - Load encoder with device_map=auto to distribute across GPUs - Only move trainable components (latent encoder/decoder, diffusion head) to GPU0 explicitly - Skip FSDP wrapping since device_map already handles distribution

- 35B MoE model can't fit on 2 GPUs when both torchrun ranks load it - Frozen encoder doesn't need gradients, so CPU is fine - Trainable components (Perceiver, diffusion head) stay on GPU - Forward pass moves encoder hidden states to GPU after encode

- Add benchmark_ldlm.py: inference throughput test (~745 tok/s on RTX 5090) - Fix dataset.py: append data_path instead of data_files list - Fix autoencoder.py: PreNorm.forward() pass context for cross-attention - Fix autoencoder.py: encode() moves tensors to CPU for encoder forward - Fix train_ldlm.py: use torch.optim.AdamW and ConstantLR directly - Update qwen3_6_27b_ldlm.yaml: local paths, reduced depth/seq_len for GPU fit

- Add move_encoder_to_gpus() for device_map='auto' encoder sharding - Fix sampler device handling for multi-GPU setups - Fix ldlm_forward device mismatches (sampler timesteps, sequence length) - Fix sequence length mismatch in reconstruction losses - Add autocast(bfloat16) for mixed-precision training - Update configs: eager attention, reduced depth, DCP checkpoint manager - Successfully training Qwen3.6-35B-A3B LDLM on RTX 5090 + RTX PRO 4000

- Sigmoid warmup schedule (Eq 29-30): gamma(s) with k=10, c=0.8 instead of linear - Self-conditioning (Section 4.2): 50% chance to feed previous z_hat estimate - Hidden state normalization: running mean/var stats for zero-mean unit-var - Tangent noise schedule (d=3): alpha_bar = 1 - t^3 from COSMOS - DiffusionHead: proj_in layer for self-conditioning input concatenation - Add warmup_gamma_min, tangent_d, self_condition_prob to configs - Return logits from ldlm_forward for wandb text logging

…ning

…pl) and training recipe

Layers the Cola DLM (arXiv:2605.06548) recipe onto the existing Repr-Align training path as an opt-in auxiliary head. LDLM (veomni/models/ldlm/) is intentionally untouched. New package veomni/models/cola_ldm/: - TextVAEEncoder: hierarchical Perceiver compressor (global + local latents, fused so locals carry global context) - BlockCausalDiT: paper-style denoiser with sinusoidal timestep embedding, AdaLN-Zero conditioning, GELU(tanh) MLP, zero-init output projection. Lifted from /home/johndpope/Documents/GitHub/ Cola-DLM (the official ByteDance release) and trimmed to drop KV-cache / NA variable-length mask / RoPE (inference-only concerns). - ColaDLMHead: encoder + DiT + loss. Supports Flow Matching (cola_prediction='v', paper default) or legacy x0-prediction. - ColaReprAlignWrapper: drop-in nn.Module bundling base LM with the head. Pass-through for FSDP/optimizer/checkpointer. Detaches student hidden states by default so the head can't destabilize Repr-Align convergence. Trainer integration (tasks/train_torch.py): - Wraps the model behind 'if train.cola_wt > 0'; off (and behaviorally identical to today's Repr-Align run) when cola_wt == 0. - Surfaces rich wandb diagnostics: losses/cola_diff, cola_pred_cosine, cola_pred_snr, per-scale latent norms/stds, cola/grad_norm, and periodic cola_hist/z_global, cola_hist/z_local histograms. Config + docs: - configs/pretrain/qwen3_6_35b_a3b_cola_ldm.yaml: 2-GPU FSDP1 launch for Qwen3.6-35B-A3B with cola_wt=0.5 defaults. - docs/cola_ldm.md: architecture, run command, metrics, extension hooks. - docs/cola_ldm_roadmap.md: validation plan (R0 vs R1 baseline), decision tree across four generation pathways (A: discard at inference, B: add decoder, C: latent conditioning, D: full Cola Text VAE port), per-pathway wandb experiments, and an inference- throughput calibration from the official Cola repo bench (37.8 tok/s aggregate at default settings on Blackwell PRO 4000). CLAUDE.md: refreshed the Architecture section to call out the two training scripts (train_torch.py vs train_ldlm.py), the two diffusion paths (Repr-Align vs LDLM), and the relevant config groups. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Repr-Align's "teacher" is a frozen snapshot of the student's init — its hidden states are deterministic and never change for the whole run, so recomputing them every step (the existing copy.deepcopy path) is pure waste. This commit lets you precompute teacher hiddens once, dump selected layers to disk, then load anchors from cache during training without ever instantiating a second model copy. New: - scripts/precompute_anchor.py — standalone dumper indexed by sha256 of input_ids; resumable; manifest captures (model, tokenizer, layers). - veomni/models/cached_teacher.py — drop-in replacement for the deepcopy'd teacher; verifies cache contract against student config. - configs/pretrain/qwen3_1_7b_repr_align_smoke.yaml — $0 end-to-end validation run for the MSI 5090. - docs/local_training.md — verified MSI + Z6 G4 inventory, memory budget for 35B-A3B, upgrade-path matrix, smoke-test procedure. Wiring: - arguments.py: train.anchor_cache_dir + train.align_layers. - auto.py / loader.py: build CachedTeacher when cache dir is set, else fall back to the existing deepcopy path. - modeling_qwen3.py: honour align_layers when computing repr_align loss (required for the cache path; backwards compatible when unset). - train_torch.py: thread the two new args through build_foundation_model. Also includes parked tweaks already in the worktree (README, qwen2 modeling tweaks, ldlm config bump). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

- train_torch.py: fall back to gloo + world_size=1 when WORLD_SIZE unset - arguments.py: default LOCAL_RANK/RANK/WORLD_SIZE to 0/0/1 - loader.py: place live teacher on cuda:1 when multi-GPU is available - module_utils.py: stray functools import Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…rs 5.x - qwen3 PreTrainedModel: add _supports_flash_attn = True (transformers 5.x renamed the dispatch attribute; the old _supports_flash_attn_2 is kept for BC error messages only and no longer used by _flash_attn_can_dispatch). - precompute_anchor.py: reproduce trainer's tokenization exactly. Old script did tok(text, padding="max_length", truncation=True, ...), one padded row per text. Trainer does encode + [eos] then split_into_chunks(tokens, msl) with no padding. Result: every cached hash was wrong, so every lookup at train time raised KeyError. New script chunks the same way and stores one cache file per chunk, keyed by sha256(input_ids[chunk]). - CachedTeacher: trainer's collator packs multiple chunks into one rmpad row with position_ids resetting per chunk. Old code hashed the whole packed row, which never matches any cached chunk hash. Now splits the row by pos2culen(position_ids), looks up each chunk, and concatenates the cached hidden states back into a packed [1, total_len, D] tensor. Added chunk length + first16 token IDs to the miss-error message so future mismatches are diagnosable without instrumentation. - smoke YAML: point at data_smoke_1000.jsonl (head -n 1000 of the 100k file) because train_size in TrainingArguments controls steps, not dataset slicing — the dataloader iterates the full JSONL and would otherwise pull uncached examples from beyond row 1000. Re-enable anchor_cache_dir, switch optimizer to anyprecision_adamw (saves ~7 GB by storing Adam states in bf16) so the run fits on the 5090. max_seq_len kept at 2048 to match the precompute run. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…r training Adds `quantize_frozen` / `quantize_frozen_dtype` to TrainingArguments. When true, every nn.Linear whose direct params are all frozen is replaced with a torchao int8 (or int4) weight-only quantized version after freeze_layers and before FSDP wrap. Cuts resident weight VRAM ~2-4x. Trainable layers stay in bf16. Smoke on Qwen3-1.7B with layers 14..27 frozen (59% of params): baseline: 6.41 GB built, 24.15 GB peak, 15.5 s/step +int8: 4.74 GB built, 13.04 GB peak, 17.2 s/step (-46% peak VRAM) loss identical within noise (16.18 vs 16.06). DCP save crashes with `cannot pickle code objects` because AffineQuantizedTensor metadata contains code refs. Fix in ModelState: dequantize quantized tensors to bf16 inside state_dict() before DCP plans the save. The live model stays quantized; only the on-disk checkpoint is bf16. Loading back into a quant model is unsupported — load plain bf16, then re-run quantize_(model, ...). Configs: - qwen3_1_7b_repr_align_quant_smoke.yaml — variant of the working smoke, freezes layers 14..27 + lm_head/embed/norm, quantize_frozen=true. - qwen3_6_27b_one_layer_repr_align.yaml — 27B dense, FSDP across 5090+PRO 4000, layer 0 trainable, anchors at 16/32/48/64. quantize_frozen off by default with a comment showing the size budget (54 GB bf16 -> 27 GB int8, fits on the 5090 alone). - qwen3_6_35b_a3b_one_layer_repr_align.yaml — 35B-A3B MoE counterpart, FSDP + CPU offload mandatory at bf16. int8 brings it to ~35 GB so the offload knob can be flipped back off. Also adds _supports_flash_attn = True to qwen3_5 PreTrainedModel for transformers 5.x dispatch compat (same fix already in qwen3; the MoE arch inherits). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Add DeepSpeed support to Open-dLLM for training models that exceed combined VRAM (e.g. Qwen3.6-35B-A3B @ 70 GB on 56 GB box): Arguments: - Add 'deepspeed' to data_parallel_mode literal - Add ds_zero_stage, ds_offload_optimizer, ds_offload_param, ds_nvme_path, ds_overlap_comm, ds_contiguous_gradients, ds_config_path - Validation: offload_param requires ZeRO-3, NVMe requires existing dir - Auto-force init_device='cpu' for deepspeed mode New files: - veomni/distributed/deepspeed_init.py: build_ds_config() + init_deepspeed_engine() via deepspeed.initialize() - veomni/checkpoint/ds_checkpointer.py: save/load via engine.save_checkpoint/load_checkpoint Distribution: - parallel_state: deepspeed mode skips DeviceMesh, dp_group=WORLD - torch_parallelize: deepspeed mode skips FSDP/DDP wrapping - checkpointer: deepspeed returns DeepSpeedCheckpointer directly Training loop (train_torch.py): - deepspeed.initialize() after optimizer + lr_scheduler built - engine.backward(loss) in backward pass - engine.step() handles grad clipping + optimizer + zero_grad - engine.get_global_grad_norm() for logging - all_reduce on dist.group.WORLD for deepspeed Deps: setup.py deepspeed extra (deepspeed>=0.15.0) Also includes pre-existing: initialize.py CPU-load support, precompute_anchor.py --max_memory arg, config pointing to local HF cache.

- torch_parallelize.py: deepspeed mode now loads weights distributed across ranks using parallel_load_safetensors (avoids each rank loading full 70 GB into CPU RAM). Materializes on CPU via NCCL broadcast, not CUDA, so DeepSpeed Engine receives a CPU model. - auto.py: all ranks load full weights for deepspeed (meta-init handled by build_parallelize_model instead) - arguments.py: allow init_device='meta' for deepspeed mode (was forced to 'cpu') - config: use init_device: meta for distributed weight loading

…jects None value)

Integrate PEFT QLoRA (4-bit NF4 + LoRA) into the Repr-Align training pipeline. Cuts peak VRAM from 31 GiB → 3.5 GiB for Qwen3-1.7B. Changes: - veomni/models/qlorafy.py — QLoRA adapter: loads model in NF4 via bitsandbytes, attaches LoRA adapters (r=16) via PEFT - veomni/models/auto.py — add enable_qlorafy path in build_foundation_model - veomni/utils/arguments.py — add enable_qlorafy + qlorafy_config fields - tasks/train_torch.py — wire QLoRA args into build_foundation_model call - configs/pretrain/smoke_qlorafy.yaml — smoke test config - pyproject.toml — add blackwell deps + CUDA 13.0 index Smoke test: 5/5 steps, loss 17.29 → 6.02, 3.5 GiB peak. Checkpoint saving has a minor issue with Params4bit serialization (separate fix). Training loop is fully functional.

…isable DCP save - qlorafy.py: set language_model_only=True for Qwen3.5 VL models (skips 4.7 GiB vision encoder that OOMs on 24GB GPUs), add max_memory={0:30GiB} to prevent BnB NF4 from rejecting CPU-offloaded modules - qlorafy_27b.yaml: set save_epochs:0 (DCP cannot serialize Params4bit) - data_loader.py/arguments.py: prefetch_factor Optional[int]=None (fixes torch DataLoader warning when num_workers=0)

- train_torch.py: log qlora/param_norm, qlora/grad_norm, qlora/grad_to_param_ratio, system/vram_allocated_gb, system/vram_reserved_gb to wandb every step - qlorafy_27b_train.yaml: training config (1000 examples, 4 layers, 2000 steps, wandb enabled, grad_accum=16)

- train_torch.py: every 100 steps, run diffusion_generate on a fixed prompt and log the output as wandb.Html for visual inspection - scripts/generation_probe.py: standalone probe for GPU 1 (requires 32GB+ GPU due to Qwen3.5 vision encoder init overhead)

…uantization, manifest output - Replace output_hidden_states=True (OOM) with forward hooks that GPU→CPU copy per layer - Add --quantize (4bit/8bit) using BitsAndBytesConfig for models that don't fit in VRAM - Add manifest.json output for CachedTeacher compatibility - Update dump-anchors.sh: single GPU with 4-bit, 4 layers (16,32,48,64), 160k ctx - Add test_anchor_smoke.sh for pipeline validation

Multi-block decoder implementing d3LLM-style pipelined parallel inference: - Block-causal attention mask (blocks attend to prompt + previous blocks + self) - Block state machine (Inactive → Activated → Fully-Activated → Completed) - Entropy-thresholded token selection with forced progress guarantee - EOS early stopping with block state cleanup - Configurable block_size, thresholds, and decoding steps Mixed into Qwen2ForCausalLM, Qwen3ForCausalLM, Qwen3_5ForCausalLM, Qwen3_5MoeForCausalLM via MultiBlockDecoderMixin. KV-cache variant deferred — requires custom attention impl for block-causal mask compatibility with HF DynamicCache. Files: +veomni/models/transformers/qwen2/multi_block_generation.py (new, ~350 LoC) ~modeling_qwen2.py, modeling_qwen3.py, modeling_qwen3_5.py, modeling_qwen3_5_moe.py ~README.md (multi-block decoder section)

…+ strategy doc

- train_torch.py: import IGNORE_INDEX; align logits/labels lengths in trajectory entropy-reg block (collator pads labels +1) - modeling_qwen3_5.py: linear-attn q/gate/beta projections use linear_key_head_dim (128) not head_dim (256) so q/k dims match and FLA chunk_gated_delta_rule stops illegal-memory-access crashing - dump-anchors.sh: match training config (seq_len, output dir, max_examples) instead of seq 160000 + unused 12TB path Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

d3LLM - smoke test

…ation https://wandb.ai/snoozie/open-dllm-compare/runs/

fixes / new charts

johndpope and others added 30 commits May 15, 2026 20:18

feat: auto-prune old checkpoints (keep latest 2)

b363683

Add prune_checkpoints() to train_ldlm.py — runs after every Checkpointer.save() on rank 0, deletes all but the 2 most recent global_step_* directories to save disk space.

fix: auto-compute decoder nhead from dim (2048/16=128)

ac56db4

- MoE model has hidden_size=2048, nhead=12 doesn't divide evenly - Auto-select nhead from [16,12,8,4,2] based on what divides dim - 2048/16=128 for Qwen3.6-35B-A3B, 5120/16=320 for 27B dense

docs: add LDLM section with throughput benchmark and training guide

b8dc423

feat: add Qwen3.6-35B-A3B LDLM benchmark (3,238 tok/s) and update docs

bdcfd58

docs: add reproducible benchmark commands and honest assumptions/caveats

fdeeeea

feat: log reconstruction text samples to wandb (target vs predicted)

4c49b69

fix: device mismatch in hidden state normalization buffers

886cdb0

fix: use native checkpoint manager, increase save_steps to 500

b6c901e

fix: reduce diffusion head depth to 3 to avoid OOM with self-conditio…

aa9c709

…ning

docs: add LDLM architecture comparison table (paper vs our Qwen3.6 im…

332e7fa

…pl) and training recipe

ok

400bd93

wip

339cf57

johndpope and others added 30 commits May 20, 2026 13:27

fix: remove ds_offload_optimizer: null from smoke config (argparse re…

a4ceb25

…jects None value)

ok

76c9a7f

updates

27e5932

add QLoRA 27B Repr-Align reproduction guide to README

0376443

Create qlorafy_27b_train.yaml

ca4be88

updates

992bada

ok

c25910d

ok

48fae6b

VFM: noise adapter + flow map wrapper + joint training objective

1303f77

perplexity

52aa0a3

VFM training script, configs, NaN skip fix, eval pipeline

f3ae467

fixes for recreation

ce42e01

qlora - replay loss

6100431

dump anchros

8837a31

wip

e922818

updates

11d9fa0

ok

f10cdd4

ok

a485e7f

fix: align QLoRA configs with anchor cache (max_seq_len=2048, rmpad) …

4bd6518

…+ strategy doc

ok

986d24c

Merge pull request #4 from scrya-com/develop

6549945

d3LLM - smoke test

Merge pull request #5 from scrya-com/feature/d3llm-trajectory-distill…

4bc352b

…ation https://wandb.ai/snoozie/open-dllm-compare/runs/

Merge pull request #6 from scrya-com/develop

af78458

fixes / new charts

Navigation Menu

WIP - updated for qwen3.6 + vastai provisioning - #27

WIP - updated for qwen3.6 + vastai provisioning#27
johndpope wants to merge 114 commits into
pengzhangzhi:mainfrom
scrya-com:main

johndpope commented May 19, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

johndpope commented May 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

johndpope commented May 19, 2026 •

edited

Loading