change configuration to v4-8 (#307)

* change configuration to v4-8

* deprecate tests for older stable diffusion models.

* linting

* deprecate tests that use sd2base checkpoint.

* update wan configs to use default 512 flash block sizes.

* fix typos and mispellings.

* revert change to scheduler test

---------

Co-authored-by: Juan Acevedo <jfacevedo@google.com>

Author: entrpn

.github/workflows/UnitTests.yml

@@ -31,7 +31,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        tpu-type: ["v5p-8"]
+        tpu-type: ["v4-8"]
     name: "TPU test (${{ matrix.tpu-type }})"
     runs-on: ["self-hosted","${{ matrix.tpu-type }}"]
     steps:

README.md

@@ -279,7 +279,7 @@ After installation completes, run the training script.
 
   ### Deploying with XPK
 
-  This assummes the user has already created an xpk cluster, installed all dependencies and the also created the dataset from the step above. For getting started with MaxDiffusion and xpk see [this guide](docs/getting_started/run_maxdiffusion_via_xpk.md).
+  This assumes the user has already created an xpk cluster, installed all dependencies and the also created the dataset from the step above. For getting started with MaxDiffusion and xpk see [this guide](docs/getting_started/run_maxdiffusion_via_xpk.md).
 
   Using v5p-256 Then the command to run on xpk is as follows:
 

docs/data_README.md

@@ -5,7 +5,7 @@ Currently MaxDiffusion supports 3 data input pipelines, controlled by the flag `
 | Pipeline | Dataset Location | Dataset formats | Features or limitations |
 | -------- | ---------------- | --------------- | ----------------------- |
 | HuggingFace (hf) | datasets in HuggingFace Hub or local/Cloud Storage | Formats supported in HF Hub: parquet, arrow, json, csv, txt | data are not loaded in memory but streamed from the saved location, good for big dataset | 
-| tf | dataset will be downaloaded form HuggingFace Hub to disk | Formats supported in HF Hub: parquet, arrow, json, csv, txt | Will read the whole dataset into memory, works for small dataset |
+| tf | dataset will be downloaded form HuggingFace Hub to disk | Formats supported in HF Hub: parquet, arrow, json, csv, txt | Will read the whole dataset into memory, works for small dataset |
 | tfrecord | local/Cloud Storage | TFRecord | data are not loaded in memory but streamed from the saved location, good for big dataset |
 | Grain | local/Cloud Storage | ArrayRecord (or any random access format) | data are not loaded in memory but streamed from the saved location, good for big dataset, supports global shuffle and data iterator checkpoint for determinism (see details in [doc](https://github.com/AI-Hypercomputer/maxtext/blob/main/getting_started/Data_Input_Pipeline.md#grain-pipeline---for-determinism)) |
 

docs/train_README.md

@@ -129,7 +129,7 @@ Now let's change the configuration as follows:
 
 Then our mesh will look like `Mesh('data': 2, 'fsdp': 2, 'tensor': 1)`.
 
-The `logical_axis_rules` specifies the sharding across the mesh. You are encouranged to add or remove rules and find what best works for you. 
+The `logical_axis_rules` specifies the sharding across the mesh. You are encouraged to add or remove rules and find what best works for you. 
 
 ### Checkpointing
 

src/maxdiffusion/configs/base_wan_14b.yml

@@ -74,13 +74,15 @@ attention_sharding_uniform: True
 dropout: 0.1
 
 flash_block_sizes: {
-  "block_q" : 2048,
+  "block_q" : 512,
   "block_kv_compute" : 512,
-  "block_kv" : 2048,
-  "block_q_dkv" : 2048,
-  "block_kv_dkv" : 2048,
+  "block_kv" : 512,
+  "block_q_dkv" : 512,
+  "block_kv_dkv" : 512,
   "block_kv_dkv_compute" : 512,
-  "use_fused_bwd_kernel": True
+  "block_q_dq" : 512,
+  "block_kv_dq" : 512,
+  "use_fused_bwd_kernel": False,
 }
 # Use on v6e
 # flash_block_sizes: {

src/maxdiffusion/configs/base_wan_27b.yml

@@ -73,14 +73,15 @@ attention_sharding_uniform: True
 dropout: 0.1
 
 flash_block_sizes: {
-  "block_q" : 1024,
-  "block_kv_compute" : 256,
-  "block_kv" : 1024,
-  "block_q_dkv" : 1024,
-  "block_kv_dkv" : 1024,
-  "block_kv_dkv_compute" : 256,
-  "block_q_dq" : 1024,
-  "block_kv_dq" : 1024
+  "block_q" : 512,
+  "block_kv_compute" : 512,
+  "block_kv" : 512,
+  "block_q_dkv" : 512,
+  "block_kv_dkv" : 512,
+  "block_kv_dkv_compute" : 512,
+  "block_q_dq" : 512,
+  "block_kv_dq" : 512,
+  "use_fused_bwd_kernel": False,
 }
 # Use on v6e
 # flash_block_sizes: {

src/maxdiffusion/configuration_utils.py

@@ -145,7 +145,7 @@ def __getattr__(self, name: str) -> Any:
     """The only reason we overwrite `getattr` here is to gracefully deprecate accessing
     config attributes directly. See https://github.com/huggingface/diffusers/pull/3129
 
-    Tihs funtion is mostly copied from PyTorch's __getattr__ overwrite:
+    Tihs function is mostly copied from PyTorch's __getattr__ overwrite:
     https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
     """
 
@@ -540,7 +540,7 @@ def extract_init_dict(cls, config_dict, **kwargs):
           f"{cls.config_name} configuration file."
       )
 
-    # 5. Give nice info if config attributes are initiliazed to default because they have not been passed
+    # 5. Give nice info if config attributes are initialized to default because they have not been passed
     passed_keys = set(init_dict.keys())
     if len(expected_keys - passed_keys) > 0:
       logger.info(f"{expected_keys - passed_keys} was not found in config. Values will be initialized to default values.")

src/maxdiffusion/generate.py

@@ -66,7 +66,7 @@ def loop_body(step, args, model, pipeline, prompt_embeds, guidance_scale, guidan
 
   # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
   # Helps solve overexposure problem when terminal SNR approaches zero.
-  # Empirical values recomended from the paper are guidance_scale=7.5 and guidance_rescale=0.7
+  # Empirical values recommended from the paper are guidance_scale=7.5 and guidance_rescale=0.7
   noise_pred = jax.lax.cond(
       guidance_rescale[0] > 0,
       lambda _: rescale_noise_cfg(noise_pred, noise_prediction_text, guidance_rescale),

src/maxdiffusion/generate_sdxl_replicated.py

@@ -32,7 +32,7 @@
 NUM_DEVICES = jax.device_count()
 
 # 1. Let's start by downloading the model and loading it into our pipeline class
-# Adhering to JAX's functional approach, the model's parameters are returned seperatetely and
+# Adhering to JAX's functional approach, the model's parameters are returned separately and
 # will have to be passed to the pipeline during inference
 pipeline, params = FlaxStableDiffusionXLPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0", revision="refs/pr/95", split_head_dim=True
@@ -83,7 +83,7 @@ def replicate_all(prompt_ids, neg_prompt_ids, seed):
 # to the function and tell JAX which are static arguments, that is, arguments that
 # are known at compile time and won't change. In our case, it is num_inference_steps,
 # height, width and return_latents.
-# Once the function is compiled, these parameters are ommited from future calls and
+# Once the function is compiled, these parameters are omitted from future calls and
 # cannot be changed without modifying the code and recompiling.
 def aot_compile(
     prompt=default_prompt,

src/maxdiffusion/input_pipeline/_tfds_data_processing.py

@@ -97,7 +97,7 @@ def _make_tfrecord_iterator(
   # Dataset cache in github runner test doesn't contain all the features since its shared, Use the default tfrecord iterator.
   # if is_training is True, loads the training dataset. If False, loads the evaluation dataset.
 
-  # checks that the dataset path is valid. In case of gcs, the existance of the dir is not checked.
+  # checks that the dataset path is valid. In case of gcs, the existence of the dir is not checked.
   is_dataset_dir_valid = "gs://" in config.dataset_save_location or os.path.isdir(config.dataset_save_location)
 
   # Determine whether to use the "cached" dataset, which requires externally