diff --git a/aiserver.py b/aiserver.py index 660dc7cc..c0d784eb 100644 --- a/aiserver.py +++ b/aiserver.py @@ -386,7 +386,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]): model.lm_head.to(breakmodel.gpu_device) model.transformer.wte.to(breakmodel.gpu_device) model.transformer.ln_f.to(breakmodel.gpu_device) - if(not model.config.rotary): + if(not hasattr(model.config, 'rotary') or not model.config.rotary): model.transformer.wpe.to(breakmodel.gpu_device) gc.collect() if(vars.bmsupported and args.breakmodel): @@ -436,7 +436,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]): model.lm_head.to(breakmodel.gpu_device) model.transformer.wte.to(breakmodel.gpu_device) model.transformer.ln_f.to(breakmodel.gpu_device) - if(not model.config.rotary): + if(not hasattr(model.config, 'rotary') or not model.config.rotary): model.transformer.wpe.to(breakmodel.gpu_device) gc.collect() if(vars.bmsupported and args.breakmodel): diff --git a/breakmodel.py b/breakmodel.py index 7417db36..8154b623 100644 --- a/breakmodel.py +++ b/breakmodel.py @@ -378,7 +378,7 @@ def new_forward( inputs_embeds[:, pos:pos+emb.shape[1]] = emb offset += emb.shape[1] - if self.rotary: + if hasattr(self, 'rotary') and self.rotary: hidden_states = inputs_embeds else: position_embeds = self.wpe(position_ids)