Modernize: VeOmni 8ca09d7 migration + LLaDA2.1/2.2 support + tests

.github/workflows/tests.yml

@@ -0,0 +1,55 @@
+name: Tests
+
+on:
+  push:
+    paths:
+      - "scripts/**"
+      - "models/**"
+      - "tasks/**"
+      - "tests/**"
+      - ".github/workflows/tests.yml"
+  pull_request:
+    paths:
+      - "scripts/**"
+      - "models/**"
+      - "tasks/**"
+      - "tests/**"
+      - ".github/workflows/tests.yml"
+  workflow_dispatch:
+
+jobs:
+  cpu-tests:
+    name: "CPU tests (VeOmni-independent)"
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          persist-credentials: false
+          # tests do NOT need the VeOmni submodule (they exercise the veomni-independent
+          # merge/split + validation logic), so we skip the (large) submodule checkout.
+          submodules: false
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"   # within VeOmni's supported range (>=3.11,<3.13)
+
+      - name: Install test deps (CPU torch only)
+        run: |
+          python -m pip install --upgrade pip
+          pip install pytest pyyaml
+          pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+      - name: Run CPU test suite
+        run: |
+          # The VeOmni smoke test and the real-checkpoint integration test skip automatically
+          # (no veomni installed; LLADA2_INTEGRATION unset). What runs here is the merge/split
+          # round-trip losslessness gate + config validation.
+          pytest tests/ -v
+
+      - name: Syntax-check changed Python
+        run: |
+          python -m py_compile \
+            scripts/moe_convertor.py \
+            models/llada2_moe/compat.py \
+            models/llada2_moe/configuration_llada2_moe.py \
+            models/llada2_moe/editing.py

.gitignore

@@ -238,3 +238,4 @@ compile_commands.json
 Cargo.lock
 
 lmms-eval
+.scratch_configs/

configs/longctx/llada2_mini_longctx_128k.yaml

@@ -1,20 +1,19 @@
 # Context-extension fine-tune: LLaDA2-mini 64k -> 128k  (YaRN 4x)
 #
-# Pre-requisite: produce the 64k checkpoint first, then:
-#
+# New (post-PR#22) nested VeOmni schema. Prereq: produce the 64k checkpoint first, then:
 #   python scripts/extend_rope_context.py \
 #       --model_path  ./configs/model_configs/llada2_mini_64k \
 #       --output_path ./configs/model_configs/llada2_mini_128k \
 #       --target_length 131072 --method yarn
-#
 # Set model_path below to the 64k fine-tuned weights, not the base weights.
 
 model:
   config_path: ./configs/model_configs/llada2_mini_128k
-  model_path:  ./output/llada2_mini_longctx_64k/hf_ckpt  # 64k fine-tuned weights
+  model_path: ./output/llada2_mini_longctx_64k/hf_ckpt   # 64k fine-tuned weights
   tokenizer_path: ./LLaDA2.0-mini
-  attn_implementation: sdpa
-  moe_implementation: fused
+  ops_implementation:
+    attn_implementation: sdpa
+    moe_implementation: fused_triton
 
 data:
   data_type: text
@@ -31,50 +30,57 @@ data:
   shuffle_buffer: 5000
   noise_range_low: 0.3
   noise_range_high: 0.8
-  num_workers: 4
-  dataloader_type: native
-  drop_last: true
+  dataloader:
+    type: native
+    num_workers: 4
+    drop_last: true
+    pin_memory: true
 
 train:
-  output_dir: ./output/llada2_mini_longctx_128k
-  data_parallel_mode: fsdp2
-  tensor_parallel_size: 1
-  ulysses_parallel_size: 1
-  expert_parallel_size: 1
-  context_parallel_size: 2          # 128k sequences need CP>=2 on most GPU configs
+  dyn_bsz: false
   global_batch_size: 4
   micro_batch_size: 1
   num_train_epochs: 1
-  rmpad: false
-  rmpad_with_pos_ids: false
   bsz_warmup_ratio: 0.0
-  dyn_bsz_margin: 0
-  dyn_bsz_buffer_size: 200
-  optimizer: adamw
-  beta1: 0.9
-  beta2: 0.999
-  lr: 1.0e-5
-  lr_min: 1.0e-6
-  lr_warmup_ratio: 0.05
-  lr_decay_style: cosine
-  lr_decay_ratio: 1.0
-  weight_decay: 0.1
-  max_grad_norm: 1.0
-  enable_mixed_precision: true
-  enable_gradient_checkpointing: true
-  enable_full_shard: true
-  enable_fsdp_offload: false
-  enable_activation_offload: false
   init_device: meta
   broadcast_model_weights_from_rank0: true
   enable_full_determinism: false
   empty_cache_steps: 100
-  ckpt_manager: dcp
-  load_checkpoint_path: ""
-  save_steps: 200
-  save_epochs: 1
-  save_hf_weights: true
+  beta1: 0.9
+  beta2: 0.999
   block_diffusion_mode: false
   same_token_labels: false
-  use_wandb: false
-  log_steps: 1
+  optimizer:
+    type: adamw
+    lr: 1.0e-5
+    lr_min: 1.0e-6
+    lr_warmup_ratio: 0.05
+    lr_decay_style: cosine
+    lr_decay_ratio: 1.0
+    weight_decay: 0.1
+    max_grad_norm: 1.0
+  accelerator:
+    tp_size: 1
+    ep_size: 1
+    pp_size: 1
+    ulysses_size: 1
+    cp_size: 2                   # 128k sequences need CP>=2 on most GPU configs
+    fsdp_config:
+      fsdp_mode: fsdp2
+      offload: false
+      mixed_precision:
+        enable: true
+    offload_config:
+      enable_activation: false
+  gradient_checkpointing:
+    enable: true
+    enable_reentrant: false
+  checkpoint:
+    output_dir: ./output/llada2_mini_longctx_128k
+    manager: dcp
+    load_path: null
+    save_steps: 200
+    save_epochs: 1
+    save_hf_weights: true
+  wandb:
+    enable: false

configs/longctx/llada2_mini_longctx_256k.yaml

@@ -1,21 +1,20 @@
 # Context-extension fine-tune: LLaDA2-mini 128k -> 256k  (LongRoPE 8x)
 #
-# Pre-requisite: produce the 128k checkpoint first, then:
-#
+# New (post-PR#22) nested VeOmni schema. Prereq: produce the 128k checkpoint first, then:
 #   python scripts/extend_rope_context.py \
 #       --model_path  ./configs/model_configs/llada2_mini_128k \
 #       --output_path ./configs/model_configs/llada2_mini_256k \
 #       --target_length 262144 --method longrope
-#
 # Set model_path below to the 128k fine-tuned weights.
 # LongRoPE bootstraps long_factor/short_factor=1.0; fine-tuning learns them.
 
 model:
   config_path: ./configs/model_configs/llada2_mini_256k
-  model_path:  ./output/llada2_mini_longctx_128k/hf_ckpt  # 128k fine-tuned weights
+  model_path: ./output/llada2_mini_longctx_128k/hf_ckpt   # 128k fine-tuned weights
   tokenizer_path: ./LLaDA2.0-mini
-  attn_implementation: sdpa
-  moe_implementation: fused
+  ops_implementation:
+    attn_implementation: sdpa
+    moe_implementation: fused_triton
 
 data:
   data_type: text
@@ -32,50 +31,57 @@ data:
   shuffle_buffer: 2000
   noise_range_low: 0.3
   noise_range_high: 0.8
-  num_workers: 4
-  dataloader_type: native
-  drop_last: true
+  dataloader:
+    type: native
+    num_workers: 4
+    drop_last: true
+    pin_memory: true
 
 train:
-  output_dir: ./output/llada2_mini_longctx_256k
-  data_parallel_mode: fsdp2
-  tensor_parallel_size: 1
-  ulysses_parallel_size: 1
-  expert_parallel_size: 1
-  context_parallel_size: 4          # 256k requires CP>=4; increase if OOM
+  dyn_bsz: false
   global_batch_size: 2
   micro_batch_size: 1
   num_train_epochs: 1
-  rmpad: false
-  rmpad_with_pos_ids: false
   bsz_warmup_ratio: 0.0
-  dyn_bsz_margin: 0
-  dyn_bsz_buffer_size: 200
-  optimizer: adamw
-  beta1: 0.9
-  beta2: 0.999
-  lr: 5.0e-6
-  lr_min: 5.0e-7
-  lr_warmup_ratio: 0.05
-  lr_decay_style: cosine
-  lr_decay_ratio: 1.0
-  weight_decay: 0.1
-  max_grad_norm: 1.0
-  enable_mixed_precision: true
-  enable_gradient_checkpointing: true
-  enable_full_shard: true
-  enable_fsdp_offload: false
-  enable_activation_offload: true   # needed for 256k activations
   init_device: meta
   broadcast_model_weights_from_rank0: true
   enable_full_determinism: false
   empty_cache_steps: 50
-  ckpt_manager: dcp
-  load_checkpoint_path: ""
-  save_steps: 100
-  save_epochs: 1
-  save_hf_weights: true
+  beta1: 0.9
+  beta2: 0.999
   block_diffusion_mode: false
   same_token_labels: false
-  use_wandb: false
-  log_steps: 1
+  optimizer:
+    type: adamw
+    lr: 5.0e-6
+    lr_min: 5.0e-7
+    lr_warmup_ratio: 0.05
+    lr_decay_style: cosine
+    lr_decay_ratio: 1.0
+    weight_decay: 0.1
+    max_grad_norm: 1.0
+  accelerator:
+    tp_size: 1
+    ep_size: 1
+    pp_size: 1
+    ulysses_size: 1
+    cp_size: 4                   # 256k requires CP>=4; increase if OOM
+    fsdp_config:
+      fsdp_mode: fsdp2
+      offload: false
+      mixed_precision:
+        enable: true
+    offload_config:
+      enable_activation: true    # needed for 256k activations
+  gradient_checkpointing:
+    enable: true
+    enable_reentrant: false
+  checkpoint:
+    output_dir: ./output/llada2_mini_longctx_256k
+    manager: dcp
+    load_path: null
+    save_steps: 100
+    save_epochs: 1
+    save_hf_weights: true
+  wandb:
+    enable: false