code: def test_lora(
      self, batch_size, model_name, zero_stage, offload_device
  ):
      local_rank = int(os.getenv("LOCAL_RANK", "0"))
      model = self.get_model(model_name)
      tokenizer = self.get_tokenizer(model_name)
      train_sentences = self.get_train_sentences(batch_size)
  
      # Inject LoRA
      model = convert_linear_layer_to_lora(model, "", 8)
      model = only_optimize_lora_parameters(model)
  
      ds_config = {
          "optimizer": {
              "type": "Adam",
              "params": {"lr": 1.0, "betas": [0.9, 0.95]},
          },
          "train_batch_size": batch_size,
          "fp16": {"enabled": True, "initial_scale_power": 12},
          "hybrid_engine": {"enabled": True, "pin_parameters": True},
          "zero_optimization": {
              "stage": zero_stage,
              "offload_optimizer": {"device": offload_device},
          },
      }
  
      model, *_ = deepspeed.initialize(model=model, config=ds_config)
  
      # Verify gradient norm is larger than 0
      before_grad_update_layer0_params = [
          ele.detach().cpu().float().numpy()
          for ele in model.layer_params[0]
          if ele is not None and len(ele.shape) > 1
      ]
  
      model.train()
      batch = tokenizer(
          train_sentences,
          max_length=16,
          padding="max_length",
          truncation=True,
          return_tensors="pt",
      )
      batch = to_device(batch, f"cuda:{local_rank}")
      batch["labels"] = batch["input_ids"]
      outputs = model(**batch, use_cache=False)
      loss = outputs.loss
      model.backward(loss)
  
      grad_norm_dict = dict()
      for name, param in model.named_parameters():
          if param.requires_grad is True:
              grad_norm_dict[name] = torch.linalg.norm(
                  safe_get_full_grad(param)
              )
  
      model.step()
      grad_norm = sum(
          [
              ele.detach().cpu().numpy()
              for ele in grad_norm_dict.values()
          ]
      )
      assert grad_norm > 1e-5
  
      # Verify parameter remains the same
      after_grad_update_layer0_params = [
          ele.detach().cpu().float().numpy()
          for ele in model.layer_params[0]
          if ele is not None and len(ele.shape) > 1
      ]
      for lhs, rhs in zip(
          before_grad_update_layer0_params,
          after_grad_update_layer0_params,
      ):
          npt.assert_allclose(lhs, rhs, 1e-5, 1e-5)
  
      # Verify fuse will mutate layer_params
      model.eval()
      with GatheredParameters(model.parameters()):
          model.fuse_lora_weight()
  
      after_grad_update_layer0_params_lora_fused = [
          ele.detach().cpu().float().numpy()
          for ele in model.layer_params[0]
          if ele is not None and len(ele.shape) > 1
      ]
  
      for lhs, rhs in zip(
          before_grad_update_layer0_params,
          after_grad_update_layer0_params_lora_fused,
      ):
          with pytest.raises(AssertionError):
              npt.assert_allclose(lhs, rhs, 1e-5, 1e-5)
  
      with GatheredParameters(model.parameters()):
          model.unfuse_lora_weight()
  
imports: import os
  import math
  import torch
  import torch.nn.functional
  import pytest
  import deepspeed
  from deepspeed.runtime.zero import GatheredParameters
  from deepspeed.ops.op_builder import OpBuilder
  from deepspeed.utils import safe_get_full_grad
  import numpy.testing
  from unit.common import DistributedTest
  from deepspeed.ops.op_builder import InferenceBuilder
  from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
label_desc: Use image node_id 4 for calling the TestHybridEngineLoRA obj's underlying member method code with example usage: obj.test_lora(batch_size, model_name, zero_stage, offload_device) without return types