Skip to content

Add native DFlash repository LoRA pilot - #1

Open
MarinaMackay wants to merge 1 commit into
mainfrom
yitong
Open

MarinaMackay wants to merge 1 commit into
mainfrom
yitong

Conversation

@MarinaMackay

Copy link
Copy Markdown
Owner

Native DFlash Repository LoRA Pilot on Two RTX 4090 GPUs

Summary

This change adds an end-to-end native DFlash training and evaluation path for
repository-specific speculative decoding on a two-RTX-4090 machine.

The implementation includes:

  • distributed Stage 1 target LoRA training and evaluation;
  • target-adapter merge and frozen-target teacher generation;
  • teacher repair, quality filtering, and ModelOpt conversation export;
  • native DFlash launch configuration with sequence length 4096, block size 8,
    512 anchors, five drafter layers, D-PACE, and self-logit distillation;
  • parameter-efficient DFlash LoRA injection, serialization, and restoration;
  • exact online greedy acceptance evaluation;
  • deterministic disjoint repository train/evaluation splitting;
  • reproducible sweep and final holdout scripts.
  • a scoped ModelOpt compatibility patch for the tested PyTorch 2.5 DDP
    environment and trusted local checkpoint resume.

Results

On a deterministic xarray 120/30 train/holdout split with zero example-ID
overlap, repository DFlash LoRA improves mean accepted draft tokens from
0.0218 to 0.1939, an 8.88x increase. The adapter trains 7,454,720
parameters while the repository-adapted target remains frozen.

The full native DFlash configuration completes on two RTX 4090 GPUs without
OOM or NaN at approximately 19.4-19.7 GiB per GPU.

Validation

  • Python compilation for all new and modified training modules
  • Shell syntax validation for all launch scripts
  • JSON parsing for both training configurations
  • Dry-run validation for Stage 1 and native DFlash launch construction
  • Remote two-GPU generic DFlash training through 1,500 steps
  • Adapter save/load evaluation on held-out repository prompts
  • Verified zero train/evaluation ID overlap

Detailed measurements and limitations are documented in
NATIVE_DFLASH_2X4090_RESULTS.md.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant