Dev/zirui/flux mlperf

.gitignore

@@ -12,4 +12,8 @@ local/
 output
 experiment
 /data/*
+!primus/backends/diffusion/data/
+!primus/backends/diffusion/data/**
+primus/backends/diffusion/data/**/__pycache__/
+primus/backends/diffusion/data/**/*.pyc
 pp_simulation_result

examples/diffusion/README.md

@@ -1,12 +1,13 @@
-# Wan Examples
+# Diffusion Examples
 
-Wan examples exercise the independent PyTorch Diffusion backend under
-`primus/backends/diffusion`. For backend details, data/checkpoint layout, and minimal
-configs, see `primus/backends/diffusion/README.md`.
+These examples exercise the independent PyTorch diffusion backend under
+`primus/backends/diffusion`, including Wan video training and FLUX text-to-image
+training. For backend details, data/checkpoint layout, and minimal configs, see
+`primus/backends/diffusion/README.md`.
 
 ## Data
 
-The default smoke-test dataset is `zirui3/tiny-video-samples` on Hugging Face:
+Wan examples use the `zirui3/tiny-video-samples` smoke-test dataset:
 
 ```bash
 huggingface-cli download zirui3/tiny-video-samples \
@@ -22,6 +23,15 @@ Expected layout:
   data/*.mp4
 ```
 
+FLUX examples use either:
+
+- precomputed encodings: a Hugging Face `datasets` directory with
+  `t5_encodings`, `clip_encodings`, `mean`, and `logvar` fields.
+- raw image-text data: a local webdataset directory, a local JSONL file, a local
+  Hugging Face dataset, or a Hugging Face dataset repo.
+
+See `primus/backends/diffusion/README.md` for checkpoint and data layout details.
+
 ## Run
 
 Set the shared `torchrun` environment first:
@@ -34,7 +44,7 @@ export MASTER_PORT=${MASTER_PORT:-29500}
 export GPUS_PER_NODE=${GPUS_PER_NODE:-8}
 ```
 
-### Pretrain
+### Wan Pretrain
 
 ```bash
 DATASET_PATH=/data/tiny-video-samples/meta.jsonl \
@@ -53,7 +63,45 @@ torchrun \
 Use `SP_SIZE=4` or `SP_SIZE=8` to enable Ulysses sequence parallelism
 when the model head count supports it.
 
-### Posttrain
+### FLUX Precomputed Pretrain
+
+```bash
+DATASET_PATH=/data/flux_precomputed \
+EMPTY_ENCODINGS_PATH=/data/flux_empty_encodings \
+MAX_STEPS=10 \
+torchrun \
+  --nnodes="$NNODES" --node_rank="$NODE_RANK" \
+  --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \
+  --nproc_per_node="$GPUS_PER_NODE" \
+  -m primus.cli.main train pretrain \
+  --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml
+```
+
+Set `PRETRAINED_PATH=/path/to/flux1-dev.safetensors` only when you want to
+initialize the DiT from existing FLUX weights. Empty `PRETRAINED_PATH` means
+random initialization.
+
+### FLUX Raw Image-Text Pretrain
+
+```bash
+DATASET=cc12m-test \
+DATASET_PATH=/data/cc12m_test \
+T5_ENCODER=google/t5-v1_1-xxl \
+CLIP_ENCODER=openai/clip-vit-large-patch14 \
+VAE_CHECKPOINT=black-forest-labs/FLUX.1-dev/ae.safetensors \
+MAX_STEPS=10 \
+torchrun \
+  --nnodes="$NNODES" --node_rank="$NODE_RANK" \
+  --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \
+  --nproc_per_node="$GPUS_PER_NODE" \
+  -m primus.cli.main train pretrain \
+  --config examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml
+```
+
+For raw training, `T5_ENCODER`, `CLIP_ENCODER`, and `VAE_CHECKPOINT` can be local
+paths or Hugging Face identifiers.
+
+### Wan Posttrain
 
 ```bash
 INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \
@@ -70,5 +118,5 @@ torchrun \
 
 The MI355X configs use Primus-style override sections such as `training`,
 `data`, `parallelism`, `optimizer`, `runtime`, and `metrics`. The diffusion
-adapter normalizes those sections into the Wan model/dataset/trainer
+adapter normalizes those sections into diffusion model/dataset/trainer
 arguments at runtime.

examples/diffusion/configs/MI355X/flux.1_dev_t2i-pretrain.yaml

@@ -0,0 +1,52 @@
+work_group: ${PRIMUS_TEAM:local}
+user_name: ${PRIMUS_USER:local}
+exp_name: ${PRIMUS_EXP_NAME:flux.1_dev_t2i-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+platform:
+  config: platform_local.yaml
+
+modules:
+  pre_trainer:
+    framework: diffusion
+    config: pre_trainer.yaml
+
+    model: flux.1_dev_t2i.yaml
+    overrides:
+      sink_level: null
+      file_sink_level: DEBUG
+      stderr_sink_level: INFO
+
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      training:
+        local_batch_size: ${LOCAL_BATCH_SIZE:1}
+        steps: ${MAX_STEPS:50}
+        num_train_epochs: 1000
+        gradient_accumulation_steps: 1
+        output_dir: ${OUTPUT_DIR:./output/flux.1_dev_t2i-pretrain}
+        save_steps: 0
+        save_strategy: ${SAVE_STRATEGY:dit_only}
+        run_name: flux.1_dev_t2i-pretrain
+
+      data:
+        dataset_path: ${DATASET_PATH:}
+        empty_encodings_path: ${EMPTY_ENCODINGS_PATH:}
+        prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1}
+        img_size: ${IMG_SIZE:256}
+
+      parallelism:
+        sp_size: 1
+        dp_replicate: ${DP_REPLICATE:1}
+
+      optimizer:
+        lr: ${LR:2.0e-4}
+        weight_decay: 0.01
+
+      runtime:
+        attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter}
+        gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false}
+        fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true}
+        report_to: none

examples/diffusion/configs/MI355X/flux.1_dev_t2i-raw-pretrain.yaml

@@ -0,0 +1,63 @@
+work_group: ${PRIMUS_TEAM:local}
+user_name: ${PRIMUS_USER:local}
+exp_name: ${PRIMUS_EXP_NAME:flux.1_dev_t2i-raw-pretrain}
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+platform:
+  config: platform_local.yaml
+
+modules:
+  pre_trainer:
+    framework: diffusion
+    config: pre_trainer.yaml
+
+    model: flux.1_dev_t2i.yaml
+    overrides:
+      sink_level: null
+      file_sink_level: DEBUG
+      stderr_sink_level: INFO
+
+      model:
+        config:
+          encoder:
+            t5_encoder: ${T5_ENCODER:google/t5-v1_1-xxl}
+            clip_encoder: ${CLIP_ENCODER:openai/clip-vit-large-patch14}
+            autoencoder: ${VAE_CHECKPOINT:black-forest-labs/FLUX.1-dev/ae.safetensors}
+            max_t5_length: ${MAX_T5_LENGTH:256}
+
+      metrics:
+        log_freq: 1
+        enable_wandb: false
+
+      training:
+        local_batch_size: ${LOCAL_BATCH_SIZE:1}
+        steps: ${MAX_STEPS:50}
+        num_train_epochs: 1000
+        gradient_accumulation_steps: 1
+        output_dir: ${OUTPUT_DIR:./output/flux.1_dev_t2i-raw-pretrain}
+        save_steps: 0
+        save_strategy: ${SAVE_STRATEGY:dit_only}
+        run_name: flux.1_dev_t2i-raw-pretrain
+
+      data:
+        dataset_type: raw
+        dataset: ${DATASET:cc12m-test}
+        dataset_format: ${DATASET_FORMAT:webdataset}
+        dataset_path: ${DATASET_PATH:}
+        prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1}
+        img_size: ${IMG_SIZE:256}
+        skip_low_resolution: false
+
+      parallelism:
+        sp_size: 1
+        dp_replicate: ${DP_REPLICATE:1}
+
+      optimizer:
+        lr: ${LR:2.0e-4}
+        weight_decay: 0.01
+
+      runtime:
+        attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter}
+        gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false}
+        fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true}
+        report_to: none