refactor some code for similarity to sd trainers.

Author: jfacevedo-google

src/maxdiffusion/checkpointing/flux_checkpointer.py

@@ -16,12 +16,10 @@
 
 from abc import ABC
 from contextlib import nullcontext
-import os
-import json
 import functools
 import jax
 import jax.numpy as jnp
-from jax.sharding import Mesh, PositionalSharding, PartitionSpec as P
+from jax.sharding import Mesh
 import orbax.checkpoint as ocp
 import grain.python as grain
 from maxdiffusion import (
@@ -35,15 +33,19 @@
 from transformers import (CLIPTokenizer, FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel, AutoTokenizer)
 
 from maxdiffusion.checkpointing.checkpointing_utils import (
-    create_orbax_checkpoint_manager,
-    load_stable_diffusion_configs,
+    create_orbax_checkpoint_manager
 )
 from maxdiffusion.models.flux.util import load_flow_model
 
 FLUX_CHECKPOINT = "FLUX_CHECKPOINT"
-_CHECKPOINT_FORMAT_DIFFUSERS = "CHECKPOINT_FORMAT_DIFFUSERS"
 _CHECKPOINT_FORMAT_ORBAX = "CHECKPOINT_FORMAT_ORBAX"
 
+FLUX_STATE_KEY = "flux_state"
+FLUX_TRANSFORMER_PARAMS_KEY = "flux_transformer_params"
+FLUX_STATE_SHARDINGS_KEY = "flux_state_shardings"
+FLUX_VAE_PARAMS_KEY = "flux_vae"
+VAE_STATE_KEY = "vae_state"
+VAE_STATE_SHARDINGS_KEY = "vae_state_shardings"
 
 class FluxCheckpointer(ABC):
 
@@ -144,67 +146,106 @@ def _set_checkpoint_format(self, checkpoint_format):
     self.checkpoint_format = checkpoint_format
 
   def save_checkpoint(self, train_step, pipeline, train_states):
+    def config_to_json(model_or_config):
+      return json.loads(model_or_config.to_json_string())
     items = {
         "config": ocp.args.JsonSave({"model_name": self.config.model_name}),
     }
 
-    items["flux_state"] = ocp.args.PyTreeSave(train_states["flux_state"])
+    items[FLUX_STATE_KEY] = ocp.args.PyTreeSave(train_states[FLUX_STATE_KEY])
 
     self.checkpoint_manager.save(train_step, args=ocp.args.Composite(**items))
 
   def load_params(self, step=None):
 
     self.checkpoint_format = _CHECKPOINT_FORMAT_ORBAX
+  
+  def load_flux_configs_from_orbax(self):
+    # TODO - load configs from orbax
+    return None
 
-  def load_checkpoint(self, step=None, scheduler_class=None):
-    clip_encoder = FlaxCLIPTextModel.from_pretrained(
-      self.config.clip_model_name_or_path, dtype=self.config.weights_dtype
-    )
-    clip_tokenizer = CLIPTokenizer.from_pretrained(
-      self.config.clip_model_name_or_path, max_length=77, use_fast=True
-    )
+  def load_diffusers_checkpoint(self):
+    flash_block_sizes = max_utils.get_flash_block_sizes(self.config)
 
-    t5_encoder = FlaxT5EncoderModel.from_pretrained(self.config.t5xxl_model_name_or_path, dtype=self.config.weights_dtype)
-    t5_tokenizer = AutoTokenizer.from_pretrained(
-      self.config.t5xxl_model_name_or_path, max_length=self.config.max_sequence_length, use_fast=True
+    if jax.device_count() == jax.local_device_count():
+      context = jax.default_device(jax.devices("cpu")[0])
+    else:
+      context = nullcontext()
+    
+    with context:
+      clip_encoder = FlaxCLIPTextModel.from_pretrained(
+        self.config.clip_model_name_or_path, dtype=self.config.weights_dtype
+      )
+      clip_tokenizer = CLIPTokenizer.from_pretrained(
+        self.config.clip_model_name_or_path,
+        max_length=77,
+        use_fast=True
+      )
+      t5_encoder = FlaxT5EncoderModel.from_pretrained(self.config.t5xxl_model_name_or_path, dtype=self.config.weights_dtype)
+      t5_tokenizer = AutoTokenizer.from_pretrained(
+        self.config.t5xxl_model_name_or_path,
+        max_length=self.config.max_sequence_length,
+        use_fast=True
+      )
+
+      vae, vae_params = FlaxAutoencoderKL.from_pretrained(
+        self.config.pretrained_model_name_or_path,
+        subfolder="vae",
+        from_pt=True,
+        use_safetensors=True,
+        dtype=self.config.weights_dtype
+      )
+
+      # loading from pretrained here causes a crash when trying to compile the model
+      # Failed to load HSACO: HIP_ERROR_NoBinaryForGpu
+      transformer = FluxTransformer2DModel.from_config(
+        self.config.pretrained_model_name_or_path,
+        subfolder="transformer",
+        mesh=self.mesh,
+        split_head_dim=self.config.split_head_dim,
+        attention_kernel=self.config.attention,
+        flash_block_sizes=flash_block_sizes,
+        dtype=self.config.activations_dtype,
+        weights_dtype=self.config.weights_dtype,
+        precision=max_utils.get_precision(self.config),
+      )
+      transformer_eval_params = transformer.init_weights(
+        rngs=self.rng, max_sequence_length=self.config.max_sequence_length, eval_only=True
+      )
+      
+      transformer_params = load_flow_model(self.config.flux_name, transformer_eval_params, "cpu")
+
+    pipeline = FluxPipeline(
+      t5_encoder,
+      clip_encoder,
+      vae,
+      t5_tokenizer,
+      clip_tokenizer,
+      transformer,
+      None,
+      dtype=self.config.activations_dtype,
+      mesh=self.mesh,
+      config=self.config,
+      rng=self.rng
     )
-    encoders_sharding = PositionalSharding(self.devices_array).replicate()
-    partial_device_put_replicated = functools.partial(max_utils.device_put_replicated, sharding=encoders_sharding)
-    clip_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), clip_encoder.params)
-    clip_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, clip_encoder.params)
-    t5_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), t5_encoder.params)
-    t5_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, t5_encoder.params)
 
+    params = {
+      FLUX_VAE_PARAMS_KEY : vae_params,
+      FLUX_TRANSFORMER_PARAMS_KEY : transformer_params
+    }
 
+    return pipeline, params
 
-    vae, vae_params = FlaxAutoencoderKL.from_pretrained(
-      self.config.pretrained_model_name_or_path, subfolder="vae", from_pt=True, use_safetensors=True, dtype="bfloat16"
-    )
+  def load_checkpoint(self, step=None, scheduler_class=None):
 
-    flash_block_sizes = max_utils.get_flash_block_sizes(self.config)
-    # loading from pretrained here causes a crash when trying to compile the model
-    # Failed to load HSACO: HIP_ERROR_NoBinaryForGpu
-    transformer = FluxTransformer2DModel.from_config(
-      self.config.pretrained_model_name_or_path,
-      subfolder="transformer",
-      mesh=self.mesh,
-      split_head_dim=self.config.split_head_dim,
-      attention_kernel=self.config.attention,
-      flash_block_sizes=flash_block_sizes,
-      dtype=self.config.activations_dtype,
-      weights_dtype=self.config.weights_dtype,
-      precision=max_utils.get_precision(self.config),
-  )
-
-    return FluxPipeline(t5_encoder,
-                        clip_encoder,
-                        vae,
-                        t5_tokenizer,
-                        clip_tokenizer,
-                        transformer,
-                        None,
-                        dtype=self.config.activations_dtype,
-                        mesh=self.mesh,
-                        config=self.config,
-                        rng=self.rng), vae_params
+    model_configs = self.load_flux_configs_from_orbax()
+
+    pipeline, params = None, {}
+
+    if model_configs:
+      print("TODO - load configs from orbax")
+    else:
+      pipeline, params = self.load_diffusers_checkpoint()
+    
+    return pipeline, params
 

src/maxdiffusion/trainers/flux_trainer.py

@@ -22,9 +22,17 @@
 import jax
 import optax
 import jax.numpy as jnp
-from jax.sharding import PartitionSpec as P
+from jax.sharding import PositionalSharding, PartitionSpec as P
 from flax.linen import partitioning as nn_partitioning
-from maxdiffusion.checkpointing.flux_checkpointer import (FluxCheckpointer, FLUX_CHECKPOINT)
+from maxdiffusion.checkpointing.flux_checkpointer import (
+  FluxCheckpointer,
+  FLUX_CHECKPOINT,
+  FLUX_TRANSFORMER_PARAMS_KEY,
+  FLUX_STATE_KEY,
+  FLUX_STATE_SHARDINGS_KEY,
+  FLUX_VAE_PARAMS_KEY,
+  VAE_STATE_KEY,
+  VAE_STATE_SHARDINGS_KEY)
 
 from maxdiffusion.input_pipeline.input_pipeline_interface import (make_data_iterator)
 
@@ -57,7 +65,7 @@ def __init__(self, config):
       raise ValueError("this script currently doesn't support training text_encoders")
 
   def post_training_steps(self, pipeline, params, train_states, msg=""):
-    imgs = pipeline(flux_params=train_states["flux_state"],
+    imgs = pipeline(flux_params=train_states[FLUX_STATE_KEY],
                     timesteps=50,
                     vae_params=train_states["vae_state"])
     imgs = np.array(imgs)
@@ -94,11 +102,21 @@ def start_training(self):
     # create train states
     train_states = {}
     state_shardings = {}
+
+    # move params to accelerator
+    encoders_sharding = PositionalSharding(self.devices_array).replicate()
+    partial_device_put_replicated = partial(max_utils.device_put_replicated, sharding=encoders_sharding)
+    pipeline.clip_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), pipeline.clip_encoder.params)
+    pipeline.clip_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, pipeline.clip_encoder.params)
+    pipeline.t5_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), pipeline.t5_encoder.params)
+    pipeline.t5_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, pipeline.t5_encoder.params)
+
+
     vae_state, vae_state_mesh_shardings = self.create_vae_state(
-        pipeline=pipeline, params=params, checkpoint_item_name="vae_state", is_training=False
+        pipeline=pipeline, params=params[FLUX_VAE_PARAMS_KEY], checkpoint_item_name=VAE_STATE_KEY, is_training=False
     )
-    train_states["vae_state"] = vae_state
-    state_shardings["vae_state_shardings"] = vae_state_mesh_shardings
+    train_states[VAE_STATE_KEY] = vae_state
+    state_shardings[VAE_STATE_SHARDINGS_KEY] = vae_state_mesh_shardings
 
     # Load dataset
     data_iterator = self.load_dataset(pipeline, params, train_states)
@@ -107,18 +125,23 @@ def start_training(self):
 
     # don't need this anymore, clear some memory.
     del pipeline.t5_encoder
+
+    # evaluate shapes
+    
     flux_state, flux_state_mesh_shardings, flux_learning_rate_scheduler = self.create_flux_state(
-        # ambiguous here, but if self.params.get("unet") doesn't exist
+        # ambiguous here, but if params=None
         # Then its 1 of 2 scenarios:
         # 1. unet state will be loaded directly from orbax
         # 2. a new unet is being trained from scratch.
         pipeline=pipeline,
         params=None, # Params are loaded inside create_flux_state
-        checkpoint_item_name="flux_state",
+        checkpoint_item_name=FLUX_STATE_KEY,
         is_training=True,
     )
-    train_states["flux_state"] = flux_state
-    state_shardings["flux_state_shardings"] = flux_state_mesh_shardings
+    flux_state = flux_state.replace(params=params[FLUX_TRANSFORMER_PARAMS_KEY])
+    flux_state = jax.device_put(flux_state, flux_state_mesh_shardings)
+    train_states[FLUX_STATE_KEY] = flux_state
+    state_shardings[FLUX_STATE_SHARDINGS_KEY] = flux_state_mesh_shardings
     #self.post_training_steps(pipeline, params, train_states, msg="before_training")
 
     # Create scheduler
@@ -320,15 +343,15 @@ def compile_train_step(self, pipeline, params, train_states, state_shardings, da
       max_logging.log("Precompiling...")
       s = time.time()
       dummy_batch = self.get_shaped_batch(self.config, pipeline)
-      p_train_step = p_train_step.lower(train_states["flux_state"], dummy_batch, train_rngs)
+      p_train_step = p_train_step.lower(train_states[FLUX_STATE_KEY], dummy_batch, train_rngs)
       p_train_step = p_train_step.compile()
       max_logging.log(f"Compile time: {(time.time() - s )}")
       return p_train_step
 
   def training_loop(self, p_train_step, pipeline, params, train_states, data_iterator, unet_learning_rate_scheduler):
 
     writer = max_utils.initialize_summary_writer(self.config)
-    flux_state = train_states["flux_state"]
+    flux_state = train_states[FLUX_STATE_KEY]
     num_model_parameters = max_utils.calculate_num_params_from_pytree(flux_state.params)
 
     max_utils.add_text_to_summary_writer("number_model_parameters", str(num_model_parameters), writer)
@@ -352,7 +375,7 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
     last_profiling_step = np.clip(
         first_profiling_step + self.config.profiler_steps - 1, first_profiling_step, self.config.max_train_steps - 1
     )
-    start_step = get_first_step(train_states["flux_state"])
+    start_step = get_first_step(train_states[FLUX_STATE_KEY])
     _, train_rngs = jax.random.split(self.rng)
     times = []
     for step in np.arange(start_step, self.config.max_train_steps):
@@ -379,7 +402,7 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
 
       if step != 0 and self.config.checkpoint_every != -1 and samples_count % self.config.checkpoint_every == 0:
         max_logging.log(f"Saving checkpoint for step {step}")
-        train_states["flux_state"] = flux_state
+        train_states[FLUX_STATE_KEY] = flux_state
         self.save_checkpoint(step, pipeline, train_states)
 
       if self.config.enable_profiler and step == last_profiling_step:
@@ -390,7 +413,7 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
           writer, local_metrics_file, running_gcs_metrics, train_metric, self.config.max_train_steps - 1, self.config
       )
 
-    train_states["flux_state"] = flux_state
+    train_states[FLUX_STATE_KEY] = flux_state
     max_logging.log(f"Average time per step: {sum(times[2:], datetime.timedelta(0)) / len(times[2:])}")
     if self.config.save_final_checkpoint:
       max_logging.log(f"Saving checkpoint for step {step}")
@@ -402,7 +425,7 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
 def _train_step(flux_state, batch, train_rng, guidance_vec, pipeline, scheduler, config):
   _, gen_dummy_rng = jax.random.split(train_rng)
   sample_rng, timestep_bias_rng, new_train_rng = jax.random.split(gen_dummy_rng, 3)
-  state_params = {"flux_state": flux_state.params}
+  state_params = {FLUX_STATE_KEY: flux_state.params}
 
   def compute_loss(state_params):
     latents = batch["pixel_values"]
@@ -424,7 +447,7 @@ def compute_loss(state_params):
     noisy_latents = pipeline.scheduler.add_noise(scheduler, latents, noise, timesteps, flux=True)
 
     model_pred = pipeline.flux.apply(
-      {"params": state_params["flux_state"]},
+      {"params": state_params[FLUX_STATE_KEY]},
       hidden_states=noisy_latents,
       img_ids=img_ids,
       encoder_hidden_states=text_embeds,
@@ -444,7 +467,7 @@ def compute_loss(state_params):
   grad_fn = jax.value_and_grad(compute_loss)
   loss, grad = grad_fn(state_params)
 
-  new_state = flux_state.apply_gradients(grads=grad["flux_state"])
+  new_state = flux_state.apply_gradients(grads=grad[FLUX_STATE_KEY])
 
   metrics = {"scalar": {"learning/loss": loss}, "scalars": {}}