| 142 | return TrainState.create(params=params, tx=optimizer, apply_fn=None) |
| 143 | |
| 144 | def init_fn(rng): |
| 145 | rng_generator = JaxRNG(rng) |
| 146 | batch = 512 |
| 147 | if FLAGS.modality == 'text': |
| 148 | params = model.init( |
| 149 | input_ids=jnp.zeros((batch, seq_length), dtype=jnp.int32), |
| 150 | position_ids=jnp.zeros((batch, seq_length), dtype=jnp.int32), |
| 151 | attention_mask=jnp.ones((batch, seq_length), dtype=jnp.int32), |
| 152 | rngs=rng_generator(llama_config.rng_keys()), |
| 153 | ) |
| 154 | elif FLAGS.modality == 'vision,text': |
| 155 | params = model.init( |
| 156 | input_ids=jnp.zeros((batch, seq_length), dtype=jnp.int32), |
| 157 | vision_masks=jnp.zeros((batch, seq_length), dtype=bool), |
| 158 | position_ids=jnp.zeros((batch, seq_length), dtype=jnp.int32), |
| 159 | attention_mask=jnp.ones((batch, seq_length), dtype=jnp.int32), |
| 160 | rngs=rng_generator(llama_config.rng_keys()), |
| 161 | ) |
| 162 | else: |
| 163 | raise ValueError(f"Unsupported modality: {FLAGS.modality}") |
| 164 | return TrainState.create(params=params, tx=optimizer, apply_fn=None) |
| 165 | |
| 166 | def train_step(train_state, rng, batch): |
| 167 | rng_generator = JaxRNG(rng) |