Transformer Model Encoder for sequence to sequence translation.
| 335 | |
| 336 | |
| 337 | class Encoder(nn.Module): |
| 338 | """Transformer Model Encoder for sequence to sequence translation.""" |
| 339 | depth: int |
| 340 | mlp_dim: Optional[int] = None # Defaults to 4x input dim |
| 341 | num_heads: int = 12 |
| 342 | dropout: float = 0.0 |
| 343 | drop_path: float = 0.0 |
| 344 | init_values: float = None |
| 345 | remat_policy: str = "none" |
| 346 | use_flash_attn: bool = False |
| 347 | scan_mlp: bool = False |
| 348 | scan_attn: bool = False |
| 349 | mlp_chunck: int = 128 |
| 350 | dtype: Optional[Dtype] = jnp.float32 |
| 351 | param_dtype: Dtype = jnp.float32 |
| 352 | mesh: Any = None |
| 353 | use_dense_general: bool = False |
| 354 | |
| 355 | @nn.compact |
| 356 | def __call__(self, x, deterministic=True): |
| 357 | out = {} |
| 358 | dpr = [float(x) for x in np.linspace(0, self.drop_path, self.depth)] # drop path decay |
| 359 | # Input Encoder |
| 360 | |
| 361 | if self.remat_policy != "none": |
| 362 | if self.remat_policy == "minimal": |
| 363 | policy = jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims |
| 364 | elif self.remat_policy == "minimal_offloaded": |
| 365 | policy = jax.checkpoint_policies.offload_dot_with_no_batch_dims(offload_src="device", |
| 366 | offload_dst="pinned_host") |
| 367 | elif self.remat_policy == "minimal_flash": |
| 368 | policy = jax.checkpoint_policies.save_from_both_policies( |
| 369 | jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims, |
| 370 | jax.checkpoint_policies.save_only_these_names( |
| 371 | "context", |
| 372 | ), |
| 373 | ) |
| 374 | else: |
| 375 | assert self.remat_policy == "full", "Remat policy needs to be on list of remat policies" |
| 376 | policy = None |
| 377 | |
| 378 | #logging.info(f"activation checkpointing {self.remat_policy}") |
| 379 | BlockLayer = remat( # pylint: disable=invalid-name |
| 380 | Encoder1DBlock, prevent_cse=True, policy=policy, static_argnums=(1,) |
| 381 | ) # "deterministic" is a static argu |
| 382 | else: |
| 383 | BlockLayer = Encoder1DBlock |
| 384 | |
| 385 | for lyr in range(self.depth): |
| 386 | block = BlockLayer( |
| 387 | name=f"encoderblock_{lyr}", |
| 388 | mlp_dim=self.mlp_dim, num_heads=self.num_heads, dropout=self.dropout, drop_path=dpr[lyr], use_flash_attn=self.use_flash_attn, |
| 389 | init_values=self.init_values, |
| 390 | scan_mlp=self.scan_mlp, |
| 391 | scan_attn=self.scan_attn, |
| 392 | mlp_chunck=self.mlp_chunck, |
| 393 | dtype=self.dtype, |
| 394 | param_dtype=self.param_dtype, |