Fix multihost_utils.process_allgather tiled=True in hf_utils.py - #4799
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request removes the 'block_diffusion' attention type and its associated configurations, refactors the email alert system to use standard SMTP instead of Pub/Sub, simplifies the training engine by delegating gradient accumulation, and removes redundant parameter synchronization in RL training. Key feedback highlights a critical bug where checkpoints at step 0 fail to restore due to a falsy check, potential SMTP connection failures on mock servers due to unconditional TLS, broken attention type comparisons resulting from the removal of type resolution, and an opportunity to log rather than silently ignore exceptions during learning rate recording.
| if not restored_step: | ||
| return None |
There was a problem hiding this comment.
Changing the check from restored_step is None to not restored_step introduces a critical bug. Since 0 is falsy in Python, if a checkpoint is restored at step 0, not restored_step will evaluate to True, causing the method to return None and fail to restore the checkpoint.
We should revert this check to explicitly check for None.
| if not restored_step: | |
| return None | |
| if restored_step is None: | |
| return None |
| with smtplib.SMTP(smtp_server, smtp_port) as server: | ||
| server.starttls() | ||
| if smtp_user and smtp_pass: | ||
| server.login(smtp_user, smtp_pass) | ||
| server.send_message(msg) |
There was a problem hiding this comment.
Calling server.starttls() unconditionally will fail with an SMTPException when connecting to a local mock SMTP server (such as one running on port 1025) that does not support TLS. This will cause the script to exit with sys.exit(1) instead of successfully sending the message or falling back.
We should only initiate TLS if authentication credentials (smtp_user and smtp_pass) are provided.
| with smtplib.SMTP(smtp_server, smtp_port) as server: | |
| server.starttls() | |
| if smtp_user and smtp_pass: | |
| server.login(smtp_user, smtp_pass) | |
| server.send_message(msg) | |
| with smtplib.SMTP(smtp_server, smtp_port) as server: | |
| if smtp_user and smtp_pass: | |
| server.starttls() | |
| server.login(smtp_user, smtp_pass) | |
| server.send_message(msg) |
| raise ValueError("causal_block_size must be positive for block-diffusion attention") | ||
| if self.attention_kernel not in ("autoselected", "dot_product", "flash"): | ||
| raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.") | ||
| self.attention_type = attention_type |
There was a problem hiding this comment.
Since AttentionType is a standard enum.Enum (not a StrEnum), comparing a string value (e.g., from the configuration) directly with AttentionType enum members will always evaluate to False (for example, "full" == AttentionType.FULL is False).
By removing _resolve_attention_type, self.attention_type can now be a string, which will break attention mask generation and attention type checks. We should ensure self.attention_type is always normalized to an AttentionType enum member.
| self.attention_type = attention_type | |
| self.attention_type = AttentionType(attention_type) if attention_type is not None else AttentionType.GLOBAL |
| self.quant = quant | ||
| self.kv_quant = kv_quant | ||
| self.attention_type = _resolve_attention_type(self.config, attention_type) | ||
| self.attention_type = attention_type |
There was a problem hiding this comment.
Since AttentionType is a standard enum.Enum (not a StrEnum), comparing a string value directly with AttentionType enum members will always evaluate to False.
By removing _resolve_attention_type, self.attention_type can now be a string, which will break attention mask generation and attention type checks. We should ensure self.attention_type is always normalized to an AttentionType enum member.
| self.attention_type = attention_type | |
| self.attention_type = AttentionType(attention_type) if attention_type is not None else AttentionType.GLOBAL |
| try: | ||
| lr = self._learning_rate_schedule(self.train_step) | ||
| self.record_metrics("learning_rate", lr) | ||
| except Exception: # pylint: disable=broad-except | ||
| pass |
There was a problem hiding this comment.
Silently catching and ignoring all exceptions with pass can hide underlying bugs in the learning rate schedule or the training step state. It is highly recommended to at least log the exception using logging.warning so that any failures are visible in the logs.
| try: | |
| lr = self._learning_rate_schedule(self.train_step) | |
| self.record_metrics("learning_rate", lr) | |
| except Exception: # pylint: disable=broad-except | |
| pass | |
| try: | |
| lr = self._learning_rate_schedule(self.train_step) | |
| self.record_metrics("learning_rate", lr) | |
| except Exception as e: # pylint: disable=broad-except | |
| logging.warning("Failed to calculate or record learning rate: %s", e) |
1d4dc21 to
9015dd6
Compare
1 participant
Fix multihost_utils.process_allgather tiled=True in hf_utils.py