From cca5b865b0883f7f6cdb02a01d4929387585d20d Mon Sep 17 00:00:00 2001 From: Gursimran Date: Thu, 2 Apr 2026 22:56:23 -0700 Subject: [PATCH] fix(engine): tree attention broke due to indentation bug in megatron-bridge-adaptation PR #1056 (#1135) --- areal/engine/megatron_engine.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index d7a219a56b..54bdb8a8b8 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -262,15 +262,15 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec, *args, **kwargs): self._check_and_apply_fp8_config() self._validate_fp8_consistency() - with self.device: - models = make_mcore_model( - hf_config=self.hf_config, - tf_config=self.tf_config, - mcore_config=self.mcore_config, - bridge=self.bridge, - bridge_type=self.bridge_cls, - is_critic=self.config.is_critic, - ) + with self.device: + models = make_mcore_model( + hf_config=self.hf_config, + tf_config=self.tf_config, + mcore_config=self.mcore_config, + bridge=self.bridge, + bridge_type=self.bridge_cls, + is_critic=self.config.is_critic, + ) self.model = _MegatronModelList(models)