Problems of Alphafold3 Inference on Jax Framework

[15yuzhuo]

Question

When I was trying to optimize the inference time of Alphafold3 on GPU under the jax framework,
I found that the sum of the timings of each module function in the inference function was not
consistent with the time testing only on this inference function. Is this due to the jax framework?

We set time for all modules. Each module time consumption is shown in table below. Obviously, the total inference
time is not consistent with time testing on each module. So we suspect that it may be caused by the jax framework.

Inference time	185.61s
Evoformer	1.14s
Diffusion	1.09s
Confidence	2.38s
Distogram	0.01s
Generate target embedding	0.86s

Thanks for your answering!

Testing the total inference time.

 def run_inference(
      self, featurised_example: features.BatchDict, rng_key: jnp.ndarray
  ) -> model.ModelResult:
    inference_time=time.time()
    result = self._model(rng_key, featurised_example)
    print(f"inference time = {time.time()-inference_time}")

result = self._model(rng_key, featurised_example) will call the _model() function. It
mainly executes the forward fuction model.Model(self._model_config)(batch) of class Model

Inference Process

  @functools.cached_property
  def _model(
      self,
  ) -> Callable[[jnp.ndarray, features.BatchDict], model.ModelResult]:
    """Loads model parameters and returns a jitted model forward pass."""

    @hk.transform
    def forward_fn(batch):

      batch_time=time.time()
      b=model.Model(self._model_config)(batch) 
      print(f"batch time = {time.time()-batch_time}")

      return b
    
    forward_time=time.time()
    a=functools.partial(
        jax.jit(forward_fn.apply, device=self._device), self.model_params
    print(f"forward time = {time.time() - forward_time}")

    return a
    )

Call function of Class Model

  def __call__(
      self, batch: features.BatchDict, key: jax.Array | None = None
  ) -> ModelResult:
    if key is None:
      key = hk.next_rng_key()

    batch = feat_batch.Batch.from_data_dict(batch)
    # generate target embedding
    embedding_module = evoformer_network.Evoformer(
        self.config.evoformer, self.global_config
    )
    target_start_time = time.time()
    target_feat = create_target_feat_embedding(
        batch=batch,
        config=embedding_module.config,
        global_config=self.global_config,
    )
    block_until_ready(target_feat)
    print(f"generate target embedding time: {time.time() - target_start_time:.2f}s")
    # Evoformer Module
    evo_start_time = time.time()
    def recycle_body(_, args):
      prev, key = args
      key, subkey = jax.random.split(key)
      embeddings = embedding_module(
          batch=batch,
          prev=prev,
          target_feat=target_feat,
          key=subkey,
      )
      embeddings['pair'] = embeddings['pair'].astype(jnp.float32)
      embeddings['single'] = embeddings['single'].astype(jnp.float32)
      return embeddings, key

    num_res = batch.num_res

    embeddings = {
        'pair': jnp.zeros(
            [num_res, num_res, self.config.evoformer.pair_channel],
            dtype=jnp.float32,
        ),
        'single': jnp.zeros(
            [num_res, self.config.evoformer.seq_channel], dtype=jnp.float32
        ),
        'target_feat': target_feat,
    }
    representation_start_time = time.time()
    if hk.running_init():
      embeddings, _ = recycle_body(None, (embeddings, key))
    else:
      num_iter = self.config.num_recycles + 1
      embeddings, _ = hk.fori_loop(0, num_iter, recycle_body, (embeddings, key))
      block_until_ready(embeddings)
    evo_end_time = time.time()
    print(f'evoformer time: {evo_end_time - evo_start_time:.2f}s')
    # Diffusion Module
    diffusion_start_time = time.time()
    samples = self._sample_diffusion(
        batch,
        embeddings,
        sample_config=self.config.heads.diffusion.eval,
    )
    block_until_ready(samples)
    print(f'diffusion time: {time.time() - diffusion_start_time:.2f}s')
    # Confidence Module
    confidence_start_time = time.time()
    confidence_output = mapping.sharded_map(
        lambda dense_atom_positions: confidence_head.ConfidenceHead(
            self.config.heads.confidence, self.global_config
        )(
            dense_atom_positions=dense_atom_positions,
            embeddings=embeddings,
            seq_mask=batch.token_features.mask,
            token_atoms_to_pseudo_beta=batch.pseudo_beta_info.token_atoms_to_pseudo_beta,
            asym_id=batch.token_features.asym_id,
        ),
        in_axes=0,
    )(samples['atom_positions'])
    block_until_ready(confidence_output)
    print(f'generate confidence time: {time.time() - confidence_start_time:.2f}s')
    # Distogram Module
    distogram_start_time = time.time()
    distogram = distogram_head.DistogramHead(
        self.config.heads.distogram, self.global_config
    )(batch, embeddings)
    block_until_ready(distogram)
    print(f'generate distogram time: {time.time() - distogram_start_time:.2f}s')

    output = {
        'diffusion_samples': samples,
        'distogram': distogram,
        **confidence_output,
    }
    if self.config.return_embeddings:
      output['single_embeddings'] = embeddings['single']
      output['pair_embeddings'] = embeddings['pair']
    return output

[weizhui]

I'm also having the same issue, waiting for a reply.

[Sikiruu]

I have just started to work on the tuning of the JAX model, and I have same questions about this. I hope this question can be answered. Thank you very much.