Merge pull request #1790 from AI-Hypercomputer:maxtext-oom

maxtext authors · maxtext authors · commit fefce810f5d7 · 2025-05-29T14:23:35.000-07:00
PiperOrigin-RevId: 764883480
diff --git a/MaxText/maxtext_utils.py b/MaxText/maxtext_utils.py
@@ -465,7 +465,7 @@ def init_initial_state(model, tx, config, is_training, key):
       {"params": key, "dropout": key, "aqt": key},
       np.ones(input_shape, dtype=jnp.int32),
       np.ones(input_shape, dtype=jnp.int32),
-      encoder_images=np.ones(image_shape, dtype=jnp.int32),
+      encoder_images=np.ones(image_shape, dtype=jnp.int32) if config.use_multimodal else None,
   )
   if is_training:
     return init_training_state(model.apply, model_vars, tx)

Original file line number	Diff line number	Diff line change
`@@ -465,7 +465,7 @@ def init_initial_state(model, tx, config, is_training, key):`
`465`	`465`	`{"params": key, "dropout": key, "aqt": key},`
`466`	`466`	`np.ones(input_shape, dtype=jnp.int32),`
`467`	`467`	`np.ones(input_shape, dtype=jnp.int32),`
`468`		`- encoder_images=np.ones(image_shape, dtype=jnp.int32),`
	`468`	`+ encoder_images=np.ones(image_shape, dtype=jnp.int32) if config.use_multimodal else None,`
`469`	`469`	`)`
`470`	`470`	`if is_training:`
`471`	`471`	`return init_training_state(model.apply, model_vars, tx)`