diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py index 9ca313f4..a07fc563 100644 --- a/src/maxdiffusion/generate_wan.py +++ b/src/maxdiffusion/generate_wan.py @@ -302,10 +302,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): videos = call_pipeline(config, pipeline, prompt, negative_prompt, num_inference_steps=warmup_steps) if isinstance(videos, tuple): videos, warmup_trace = videos - max_logging.log( - "Warmup breakdown: " - + ", ".join(f"{stage}={seconds:.1f}s" for stage, seconds in warmup_trace.items()) - ) + max_logging.log("Warmup breakdown: " + ", ".join(f"{stage}={seconds:.1f}s" for stage, seconds in warmup_trace.items())) max_logging.log("===================== Model details =======================") max_logging.log(f"model name: {config.model_name}") diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 0ddbed62..bdb738cf 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -1113,9 +1113,7 @@ def wrap_ulysses_ring_attention(query, key, value): use_fixed_m=use_fixed_m, ) if use_fixed_m: - attention_output = jnp.swapaxes( - jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr), 2, 3 - ) + attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr), 2, 3) else: attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value), 2, 3) else: diff --git a/src/maxdiffusion/models/wan/wan_utils.py b/src/maxdiffusion/models/wan/wan_utils.py index 2904f15d..883ef7d2 100644 --- a/src/maxdiffusion/models/wan/wan_utils.py +++ b/src/maxdiffusion/models/wan/wan_utils.py @@ -383,9 +383,7 @@ def convert_chunk(ckpt_shard_path, chunk_keys): # rename_key_and_reshape_tensor only reindexes/transposes views; the # single real copy happens on assignment into the target buffer below. - flax_key, flax_tensor = rename_key_and_reshape_tensor( - pt_tuple_key, tensor, random_flax_state_dict, scan_layers - ) + flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, scan_layers) flax_key = rename_for_nnx(flax_key) flax_key = _tuple_str_to_int(flax_key) @@ -428,9 +426,7 @@ def convert_chunk(ckpt_shard_path, chunk_keys): validate_flax_state_dict(eval_shapes, flax_state_dict) flax_state_dict = unflatten_dict(flax_state_dict) - max_logging.log( - f"Converted {subfolder or 'transformer'} weights to host arrays in {time.perf_counter() - t_start:.1f}s" - ) + max_logging.log(f"Converted {subfolder or 'transformer'} weights to host arrays in {time.perf_counter() - t_start:.1f}s") return flax_state_dict diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py index d56171ff..730d67fc 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline.py @@ -279,15 +279,11 @@ def stages_via_ici(val, sharding) -> bool: put_arrays = [None] * len(put_specs) if staged_ids: - staged = jax.device_put( - [put_specs[i][1] for i in staged_ids], [dim0_sharding] * len(staged_ids) - ) + staged = jax.device_put([put_specs[i][1] for i in staged_ids], [dim0_sharding] * len(staged_ids)) # out_shardings must be the exact target sharding objects (not an # equivalent P()): downstream jit cache keys include arg shardings, so # a different-but-equivalent spec would force a full recompile. - replicate_fn = jax.jit( - lambda xs: xs, out_shardings=[put_specs[i][2] for i in staged_ids] - ) + replicate_fn = jax.jit(lambda xs: xs, out_shardings=[put_specs[i][2] for i in staged_ids]) for i, replicated in zip(staged_ids, replicate_fn(staged)): put_arrays[i] = replicated if direct_ids: diff --git a/src/maxdiffusion/tests/custom_splash_fixed_m_test.py b/src/maxdiffusion/tests/custom_splash_fixed_m_test.py index 5571007d..688d34e5 100644 --- a/src/maxdiffusion/tests/custom_splash_fixed_m_test.py +++ b/src/maxdiffusion/tests/custom_splash_fixed_m_test.py @@ -47,13 +47,9 @@ class CustomSplashFixedMTest(unittest.TestCase): def setUp(self): super().setUp() self.scale = 1.0 / math.sqrt(self.head_dim) - self.block_sizes = custom_splash._BlockSizes( - block_q=2048, block_kv=1024, block_kv_compute=512 - ) + self.block_sizes = custom_splash._BlockSizes(block_q=2048, block_kv=1024, block_kv_compute=512) - def _random_qkv( - self, q_gain: float = 1.0, k_gain: float = 1.0 - ) -> tuple[jax.Array, jax.Array, jax.Array]: + def _random_qkv(self, q_gain: float = 1.0, k_gain: float = 1.0) -> tuple[jax.Array, jax.Array, jax.Array]: """Returns bf16 (q, k, v), optionally amplifying head 0 of q and k.""" shape = (self.num_heads, self.seq_len, self.head_dim) q = jax.random.normal(jax.random.PRNGKey(0), shape, jnp.bfloat16) @@ -63,18 +59,14 @@ def _random_qkv( k = k.at[0].multiply(k_gain) return q, k, v - def _reference( - self, q: jax.Array, k: jax.Array, v: jax.Array - ) -> jax.Array: + def _reference(self, q: jax.Array, k: jax.Array, v: jax.Array) -> jax.Array: """Per-head f32 softmax attention reference.""" qf, kf, vf = (x.astype(jnp.float32) for x in (q, k, v)) logits = jnp.einsum("hsd,htd->hst", qf, kf) * self.scale probs = jax.nn.softmax(logits, axis=-1) return jnp.einsum("hst,htd->hsd", probs, vf) - def _run_kernel( - self, q: jax.Array, k: jax.Array, v: jax.Array, use_fixed_m: bool - ) -> tuple[jax.Array, jax.Array | None]: + def _run_kernel(self, q: jax.Array, k: jax.Array, v: jax.Array, use_fixed_m: bool) -> tuple[jax.Array, jax.Array | None]: """Runs the custom kernel using the production scaling convention. Args: @@ -95,9 +87,7 @@ def _run_kernel( k_in = k_in - jnp.mean(k_in, axis=1, keepdims=True) qn = jnp.sqrt((q_in.astype(jnp.float32) ** 2).sum(-1)).max(axis=1) mk_h = jnp.sqrt((k_in.astype(jnp.float32) ** 2).sum(-1)).max(axis=1) - eligible = (qn * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype( - jnp.float32 - ) + eligible = (qn * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) mk = jnp.stack([mk_h, eligible]) kernel = custom_splash.make_splash_mha( block_sizes=self.block_sizes,