add code for alignment

Author: zhaoyinglia

flagscale/models/megatron/bagel/models/bagel_model.py

@@ -377,10 +377,13 @@ def freeze(
         """
         modules = []
         if freeze_language_model and self.language_model:
+            self.language_model.eval()
             modules.append(self.language_model)
         if freeze_vision_model and self.vision_model is not None:
+            self.vision_model.eval()
             modules.append(self.vision_model)
         if freeze_vae_model and self.vae_model is not None:
+            self.vae_model.eval()
             modules.append(self.vae_model)
         if freeze_text_embed and self.language_model:
             modules.append(self.language_model.embedding.word_embeddings)
@@ -666,6 +669,10 @@ def forward(
         vit_token_seqlens: Optional[Tensor] = None,
         # VAE generation inputs
         padded_images: Optional[Tensor] = None,
+        # padded_latent: Optional[Tensor] = None,
+        # packed_latent_clean: Optional[Tensor] = None,
+        # packed_latent_noise: Optional[Tensor] = None,
+        # packed_timesteps_effective: Optional[Tensor] = None,
         patchified_vae_latent_shapes: Optional[List] = None,
         packed_latent_position_ids: Optional[Tensor] = None,
         packed_vae_token_indexes: Optional[Tensor] = None,
@@ -702,16 +709,18 @@ def forward(
                 f"but got {tuple(packed_text_ids.shape)}"
             )
 
+        # print(f"{packed_text_ids.shape=}, {torch.sum(packed_text_ids, dtype=torch.float32)=}, {packed_text_ids=}")
         text_embeddings = self.language_model.embedding(
             input_ids=packed_text_ids, position_ids=packed_position_ids.unsqueeze(0),
         )
-        print(f"{packed_text_ids.shape=}, {text_embeddings.shape=}")
+        # print(f"{packed_text_ids.shape=}, {text_embeddings.shape=}")
 
         # Create the combined sequence tensor
         # For packed sequences: [total_seq, 1, hidden]
         combined_embeddings = text_embeddings.new_zeros(size=(sequence_length, self.language_model.config.hidden_size))
-        print(f"{packed_text_indexes.shape=}, {combined_embeddings.shape=}")
+        # print(f"{packed_text_indexes.shape=}, {combined_embeddings.shape=}")
         combined_embeddings[packed_text_indexes] = text_embeddings.squeeze(1)
+        # print(f"{text_embeddings.shape=}, {torch.sum(text_embeddings, dtype=torch.float32)=}, {text_embeddings=}")
 
         # --- Step 2: Process ViT tokens (understanding path) ---
         has_vision_path = (
@@ -733,18 +742,20 @@ def forward(
             cu_seqlens = cu_seqlens.to(torch.int32)
             max_seqlen = torch.max(vit_token_seqlens).item()
             # [num_patches, 1, vit_hidden]
-            print(f"{packed_vit_tokens.shape=}, {packed_vit_tokens.shape=}")
+            # print(f"{packed_vit_tokens.shape=}, {packed_vit_tokens.shape=}")
             image_embeddings = self._embed_vision_tokens(
                 pixel_values=packed_vit_tokens,
                 position_ids=packed_vit_position_ids,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
             )
-            print(f"{image_embeddings.shape=}")
+            # print(f"after vit {image_embeddings.shape=}, {torch.sum(image_embeddings, dtype=torch.float32)=}, {image_embeddings=}")
+            # print(f"{image_embeddings.shape=}")
 
             if image_embeddings.dtype != combined_embeddings.dtype:
                 image_embeddings = image_embeddings.to(combined_embeddings.dtype)
             combined_embeddings[packed_vit_token_indexes] = image_embeddings.squeeze(1)
+            # print(f"{image_embeddings.shape=}, {torch.sum(image_embeddings, dtype=torch.float32)=}, {image_embeddings=}")
         elif has_vision_path:
             # Keep trainable vision modules in the autograd graph when this
             # microbatch contains no vision-understanding task.
@@ -771,6 +782,8 @@ def forward(
 
         # --- Step 3: Process VAE latents (generation path) ---
         packed_latent_clean = None
+        packed_latent_noise = None
+        packed_timesteps_effective = None
         noise = None
         processed_timesteps = None
         has_generation_path = (
@@ -783,26 +796,52 @@ def forward(
         )
         has_gen_tokens = (
             has_generation_path
-            and padded_images is not None
+            and (
+                packed_latent_clean is not None
+                or padded_latent is not None
+                or padded_images is not None
+            )
             and packed_vae_token_indexes is not None
             and packed_vae_token_indexes.numel() > 0
             and patchified_vae_latent_shapes is not None
             and packed_latent_position_ids is not None
             and packed_timesteps is not None
         )
         if has_gen_tokens:
-            vae_dtype = next(self.vae_model.parameters()).dtype
-            padded_latent = self.vae_model.encode(padded_images.to(vae_dtype))
-            packed_latent_clean = self._patchify_vae_latents(
-                padded_latent,
-                patchified_vae_latent_shapes,
-            )
-            packed_latent, noise, processed_timesteps = (
-                self._apply_flow_matching_noise(
-                    packed_latent_clean,
-                    packed_timesteps,
+            if packed_latent_clean is None:
+                if padded_latent is None:
+                    vae_dtype = next(self.vae_model.parameters()).dtype
+                    padded_latent = self.vae_model.encode(padded_images.to(vae_dtype))
+                packed_latent_clean = self._patchify_vae_latents(
+                    padded_latent,
+                    patchified_vae_latent_shapes,
+                )
+            if packed_latent_noise is not None:
+                noise = packed_latent_noise.to(packed_latent_clean.dtype)
+                if packed_timesteps_effective is None:
+                    processed_timesteps = torch.sigmoid(
+                        packed_timesteps.to(packed_latent_clean.dtype)
+                    )
+                    processed_timesteps = (
+                        self.timestep_shift
+                        * processed_timesteps
+                        / (1 + (self.timestep_shift - 1) * processed_timesteps)
+                    )
+                else:
+                    processed_timesteps = packed_timesteps_effective.to(
+                        packed_latent_clean.dtype
+                    )
+                packed_latent = (
+                    (1 - processed_timesteps[:, None]) * packed_latent_clean
+                    + processed_timesteps[:, None] * noise
+                )
+            else:
+                packed_latent, noise, processed_timesteps = (
+                    self._apply_flow_matching_noise(
+                        packed_latent_clean,
+                        packed_timesteps,
+                    )
                 )
-            )
             packed_latent = self._embed_generation_tokens(
                 packed_latent,
                 processed_timesteps,
@@ -811,6 +850,7 @@ def forward(
             combined_embeddings[packed_vae_token_indexes] = packed_latent.to(
                 combined_embeddings.dtype
             )
+            # print(f"{packed_latent.shape=}, {torch.sum(packed_latent, dtype=torch.float32)=}, {packed_latent=}")
         elif has_generation_path:
             # Match the real generation path with one dummy image so the VAE
             # encoder and all generation-input modules stay in the graph.
@@ -875,7 +915,7 @@ def forward(
             attn_modes=attn_modes,
             device=combined_embeddings.device,
         )
-        print(f"{packed_und_token_indexes.shape=}")
+        # print(f"{packed_und_token_indexes.shape=}")
 
         # Run through GPTModel with decoder_input (skips embedding)
         hidden_states = self.language_model.decoder(

flagscale/models/megatron/bagel/models/siglip_model.py

@@ -155,9 +155,9 @@ def forward(
             x = x + self.position_embeddings(packed_flattened_position_ids)
 
         # Reshape for TransformerBlock: [seq_len, batch=1, hidden_size]
-        x = x.unsqueeze(1)
+        x = x.unsqueeze(1).contiguous()
 
-        print(f"{cu_seqlens=}, {max_seqlen=}")
+        # print(f"{cu_seqlens=}, {max_seqlen=}")
         # Build PackedSeqParams for variable-length attention
         packed_seq_params = PackedSeqParams(
             cu_seqlens_q=cu_seqlens,
@@ -167,8 +167,9 @@ def forward(
             max_seqlen_kv=max_seqlen,
         )
 
-        print(f"{x.shape=}")
+        # print(f"{x.shape=}")
         # Transformer forward
+        # print(f"before encoder: {x.shape=}, {torch.sum(x, dtype=torch.float32)=}, {x=}, {packed_seq_params=}")
         x = self.decoder(
             hidden_states=x,
             attention_mask=None,

flagscale/train/megatron/train_bagel.py

@@ -1,6 +1,8 @@
 # Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
 """Pretrain BAGEL multimodal model for unified image understanding and generation."""
 
+import os
+
 import torch
 from megatron.core import parallel_state
 from megatron.core.enums import ModelType
@@ -129,6 +131,24 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
     return dataset_provider(train_val_test_num_samples)
 
 
+def _load_alignment_batch(path):
+    """Load and validate a CPU-only batch dumped by original BAGEL."""
+    data = torch.load(path, map_location="cpu", weights_only=False)
+    required = {
+        "sequence_length", "sample_lens", "split_lens", "attn_modes",
+        "packed_text_ids", "packed_text_indexes", "packed_position_ids",
+        "packed_latent_clean", "packed_latent_noise",
+        "packed_timesteps_effective",
+    }
+    missing = sorted(required.difference(data))
+    if missing:
+        raise ValueError(f"alignment batch is missing required fields: {missing}")
+    metadata = data.get("_alignment_metadata", {})
+    if metadata.get("source") != "original_bagel":
+        raise ValueError(f"unexpected alignment batch source: {metadata!r}")
+    return data
+
+
 def get_batch(data_iterator):
     """Generate a batch from data iterator.
 
@@ -151,8 +171,12 @@ def get_batch(data_iterator):
     assert (getattr(args, 'pipeline_model_parallel_size', 1) == 1)
     assert (getattr(args, 'tensor_model_parallel_size', 1) == 1)
 
-    data = next(data_iterator)
-    print(f"{data=}")
+    alignment_batch_path = os.getenv("BAGEL_ALIGN_BATCH")
+    if alignment_batch_path:
+        data = _load_alignment_batch(alignment_batch_path)
+    else:
+        data = next(data_iterator)
+    # print(f"{data=}")
 
     # --- Assemble batch dict ---
     batch = {
@@ -179,6 +203,14 @@ def get_batch(data_iterator):
     # --- VAE image generation fields ---
     if 'padded_images' in data:
         batch['padded_images'] = data['padded_images'].cuda(non_blocking=True)
+    if 'padded_latent' in data:
+        batch['padded_latent'] = data['padded_latent'].cuda(non_blocking=True)
+    if 'packed_latent_clean' in data:
+        batch['packed_latent_clean'] = data['packed_latent_clean'].cuda(non_blocking=True)
+    if 'packed_latent_noise' in data:
+        batch['packed_latent_noise'] = data['packed_latent_noise'].cuda(non_blocking=True)
+    if 'packed_timesteps_effective' in data:
+        batch['packed_timesteps_effective'] = data['packed_timesteps_effective'].cuda(non_blocking=True)
     if 'patchified_vae_latent_shapes' in data:
         batch['patchified_vae_latent_shapes'] = data['patchified_vae_latent_shapes']
     if 'packed_latent_position_ids' in data:
@@ -202,7 +234,7 @@ def get_batch(data_iterator):
     if 'mse_loss_indexes' in data:
         batch['mse_loss_indexes'] = data['mse_loss_indexes'].cuda(non_blocking=True)
 
-    print(f"{batch=}")
+    # print(f"{batch=}")
     return batch
 
 
@@ -263,6 +295,10 @@ def forward_step(data_iterator, model: BagelModel):
         vit_token_seqlens=batch.get('vit_token_seqlens'),
         # VAE generation
         padded_images=batch.get('padded_images'),
+        padded_latent=batch.get('padded_latent'),
+        packed_latent_clean=batch.get('packed_latent_clean'),
+        packed_latent_noise=batch.get('packed_latent_noise'),
+        packed_timesteps_effective=batch.get('packed_timesteps_effective'),
         patchified_vae_latent_shapes=batch.get('patchified_vae_latent_shapes'),
         packed_latent_position_ids=batch.get('packed_latent_position_ids'),
         packed_vae_token_indexes=batch.get('packed_vae_token_indexes'),
@@ -319,20 +355,20 @@ def forward_step(data_iterator, model: BagelModel):
         mse = mse.mean(dim=-1).sum() * dp_world_size / total_mse_tokens
         combined_loss = combined_loss + mse * getattr(args, 'mse_weight', 1.0)
 
-    print(f"{ce=}, {mse=}")
+    # print(f"{ce=}, {mse=}")
 
     combined_loss = combined_loss + dummy
 
-    if torch.distributed.get_rank() == 0:
-        print(
-            "[BAGEL_LOSS_GRAPH]",
-            f"ce_shape={None if ce is None else tuple(ce.shape)}",
-            f"ce_grad_fn={None if ce is None else ce.grad_fn}",
-            f"mse_shape={None if mse is None else tuple(mse.shape)}",
-            f"mse_grad_fn={None if mse is None else mse.grad_fn}",
-            f"combined_grad_fn={combined_loss.grad_fn}",
-            flush=True,
-        )
+    # if torch.distributed.get_rank() == 0:
+    #     print(
+    #         "[BAGEL_LOSS_GRAPH]",
+    #         f"ce_shape={None if ce is None else tuple(ce.shape)}",
+    #         f"ce_grad_fn={None if ce is None else ce.grad_fn}",
+    #         f"mse_shape={None if mse is None else tuple(mse.shape)}",
+    #         f"mse_grad_fn={None if mse is None else mse.grad_fn}",
+    #         f"combined_grad_fn={combined_loss.grad_fn}",
+    #         flush=True,
+    #     )
 
     return combined_loss.unsqueeze(0), bagel_loss_func
 
@@ -411,6 +447,22 @@ def add_bagel_extra_args(parser):
         help='Reweight CE loss using per-token ce_loss_weights.',
     )
 
+    # Random preprocessing controls. Keep these explicit so alignment runs do
+    # not silently fall back to BagelDataConfig defaults.
+    group.add_argument(
+        '--text-cond-dropout-prob', type=float, default=0.0,
+        help='Probability of dropping text conditioning during packing.',
+    )
+    group.add_argument(
+        '--vit-cond-dropout-prob', type=float, default=0.0,
+        help='Probability of dropping ViT conditioning during packing.',
+    )
+    group.add_argument(
+        '--vae-cond-dropout-prob', type=float, default=0.0,
+        help='Probability of dropping VAE conditioning during packing.',
+    )
+
+
     # Freeze controls
     group.add_argument(
         '--freeze-VAE', action='store_true', default=False,