Merge pull request #3 from aws-neuron/release_2.21.0

Neuron 2.21 release

Author: FThompsonAWS

examples/dog.jpg

@@ -0,0 +1,145 @@
+import torch
+import os
+
+from transformers import AutoTokenizer, GenerationConfig
+
+from neuronx_distributed_inference.models.config import MultimodalVisionNeuronConfig, OnDeviceSamplingConfig
+from neuronx_distributed_inference.models.mllama.modeling_mllama import MllamaInferenceConfig, NeuronMllamaForCausalLM
+from neuronx_distributed_inference.utils.hf_adapter import load_pretrained_config, HuggingFaceGenerationAdapter
+from neuronx_distributed_inference.models.mllama.model_wrapper_mllama import NUM_IMAGE_PER_PROMPT
+from neuronx_distributed_inference.models.mllama.utils import create_vision_mask, get_image, get_image_tensors, add_instruct
+from neuronx_distributed_inference.modules.generation.sampling import prepare_sampling_params
+from neuronx_distributed_inference.utils.benchmark import benchmark_sampling
+
+# TODO : Either read from os_environment var or from arg_parser.
+checkpoint = "meta"
+model_variant = "11B"
+model_path = f"/home/ubuntu/models/Llama-3.2-{model_variant}-Vision-Instruct-{checkpoint}/"
+traced_model_path = f"/home/ubuntu/workplace/traced_models/Llama-3.2-{model_variant}-Vision-Instruct-{checkpoint}/"
+
+torch.manual_seed(0)
+
+
+def run_llama_generate():
+    # Initialize configs and tokenizer.
+    batch_size = 1
+    num_img_per_prompt = 1
+    max_context_length = 1024
+    seq_len = 2048
+
+    generation_config = GenerationConfig.from_pretrained(model_path)
+    generation_config_kwargs = {
+        "top_k": 1,
+    }
+    generation_config.update(**generation_config_kwargs)
+
+    on_device_sampling_config=OnDeviceSamplingConfig(
+                                                     dynamic=True, 
+                                                     )
+
+    neuron_config = MultimodalVisionNeuronConfig(
+        tp_degree=32,
+        batch_size=batch_size,
+        max_context_length=max_context_length,
+        seq_len=seq_len,
+        on_device_sampling_config=on_device_sampling_config,
+        enable_bucketing=True,
+        sequence_parallel_enabled=False,
+        fused_qkv=False,
+        async_mode=False,
+    )
+    config = MllamaInferenceConfig(
+        neuron_config,
+        load_config=load_pretrained_config(model_path),
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_path, padding_side="right")
+    tokenizer.pad_token = tokenizer.eos_token
+
+ 
+    # Generate outputs.
+    image = get_image("dog.jpg")
+    batch_image = [[image] * num_img_per_prompt] * batch_size
+    pixel_values, aspect_ratios, num_chunks, has_image = get_image_tensors(config, batch_image)
+
+    prompt = add_instruct("What is in this image? Tell me a story", has_image)
+    batch_prompt = [prompt] * batch_size
+
+    if not os.path.exists(traced_model_path):
+        # Compile and save model.
+        print("\nCompiling and saving model...")
+        model = NeuronMllamaForCausalLM(model_path, config)
+        model.compile(traced_model_path)
+        tokenizer.save_pretrained(traced_model_path)
+
+    # Load from compiled checkpoint.
+    print("\nLoading model from compiled checkpoint...")
+    model = NeuronMllamaForCausalLM(traced_model_path)
+    model.load(traced_model_path)
+    tokenizer = AutoTokenizer.from_pretrained(traced_model_path)
+
+    print("\nGenerating outputs...")
+    print(f"Prompts: {batch_prompt}")
+
+    inputs = tokenizer(batch_prompt, padding=True, return_tensors="pt", add_special_tokens=False)
+
+    vision_token_id = tokenizer("<|image|>", add_special_tokens=False).input_ids[0]
+    vision_mask = create_vision_mask(inputs.input_ids, vision_token_id)
+
+    generation_model = HuggingFaceGenerationAdapter(model)
+
+    # Test Sampling Parameters
+    sampling_params = prepare_sampling_params(batch_size=batch_size, top_k=[1], top_p=[1.0],  temperature=[1.0])
+    outputs = generation_model.generate(
+        inputs.input_ids,
+        generation_config=generation_config,
+        attention_mask=inputs.attention_mask,
+        max_length=model.config.neuron_config.max_length,
+        sampling_params=sampling_params, 
+        pixel_values=pixel_values, 
+        aspect_ratios=aspect_ratios,
+        vision_mask =vision_mask,
+        num_chunks=num_chunks, 
+        has_image=has_image,
+        max_new_tokens=512,
+    )
+    output_tokens = tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    
+    print("Generated outputs:")
+    for i, output_token in enumerate(output_tokens):
+        print(f"Output {i}: {output_token}")
+
+
+    # Test with text-only input
+    pixel_values, aspect_ratios, num_chunks, has_image = get_image_tensors(config, [[]] * batch_size)
+
+    prompt = add_instruct("what is the recipe of mayonnaise in two sentences?", has_image)
+    batch_prompt = [prompt] * batch_size
+    inputs = tokenizer(batch_prompt, padding=True, return_tensors="pt")
+
+    sampling_params = prepare_sampling_params(batch_size=batch_size, top_k=[1], top_p=[1.0],  temperature=[1.0])
+    outputs = generation_model.generate(
+        inputs.input_ids,
+        generation_config=generation_config,
+        attention_mask=inputs.attention_mask,
+        max_length=model.config.neuron_config.max_length,
+        sampling_params=sampling_params, 
+        pixel_values=pixel_values, 
+        aspect_ratios=aspect_ratios,
+        vision_mask=vision_mask,
+        num_chunks=num_chunks, 
+        has_image=has_image,
+        max_new_tokens=512,
+    )
+    output_tokens = tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+
+    print("Generated outputs:")
+    for i, output_token in enumerate(output_tokens):
+        print(f"Output {i}: {output_token}")
+        
+    print("\nPerformance Benchmarking!")
+    benchmark_sampling(model=model, draft_model=None, generation_config=generation_config, target="all", image=True)
+
+if __name__ == "__main__":
+    run_llama_generate()
+

examples/generation_mllama.py

@@ -1,4 +1,7 @@
 transformers==4.45.*
 sentencepiece
 pillow
-pytest-forked
+pytest-forked
+tiktoken
+blobfile
+torchvision

examples/requirements.txt

@@ -0,0 +1,114 @@
+import torch
+import torch.nn as nn
+
+from neuronx_distributed_inference.models.config import InferenceConfig
+from neuronx_distributed_inference.models.mllama.modeling_mllama_vision import VisionEncoder
+from neuronx_distributed_inference.models.mllama.utils import META_CHECKPOINT, to_2tuple
+
+from .test_utils import load_checkpoint, logger, save_checkpoint, setup_debug_env, trace_nxd_model
+
+VISION_SEQ_LEN = 1601
+VISION_HIDDEN_DIM = 1280
+MAX_NUM_CHUNKS = 4
+TORCH_DTYPE = torch.float32
+
+
+class VisionEncoderPosEmbedOnly(VisionEncoder):
+    def __init__(self, max_num_tiles, image_size, patch_size, width):
+        nn.Module.__init__(self)
+        self.config = InferenceConfig(neuron_config=None)
+        self.config.checkpoint = META_CHECKPOINT
+        self.max_num_tiles = max_num_tiles
+        self.image_size = to_2tuple(image_size)
+        self.patch_size = to_2tuple(patch_size)
+        self.grid_size = (
+            self.image_size[0] // self.patch_size[0],
+            self.image_size[1] // self.patch_size[1],
+        )
+        scale = width**-0.5
+        self.positional_embedding = nn.Parameter(
+            scale * torch.randn(self.grid_size[0] * self.grid_size[1] + 1, width, dtype=TORCH_DTYPE)
+        )
+        self.gated_positional_embedding = nn.Parameter(
+            scale
+            * torch.randn(
+                max_num_tiles,
+                max_num_tiles,
+                self.grid_size[0] * self.grid_size[1] + 1,
+                width,
+                dtype=TORCH_DTYPE,
+            )
+        )
+        # Don't initialize to zero, otherwise the gated_positional_embedding has no effect on output
+        self.gated_positional_embedding_gate = nn.Parameter(torch.randn(1, dtype=TORCH_DTYPE))
+
+    def forward(self, x, ar):
+        return self.apply_positional_embedding(x, ar, ar_ids=None)
+
+
+class VisionEncoderMeta(VisionEncoderPosEmbedOnly):
+    def apply_positional_embedding(self, x, ar, ar_ids=None):
+        # apply regular position embedding
+        bsz, num_chunks, num_tokens, dim = x.shape
+        x = x.view(bsz * num_chunks, num_tokens, dim)
+        x = x + self.positional_embedding * (1 - self.gated_positional_embedding_gate.tanh())
+        x = x.view(bsz, num_chunks, num_tokens, dim)
+        for idx, arx in enumerate(ar):
+            _pos_embed = self.gated_positional_embedding[: arx[0], : arx[1]]
+            _pos_embed = _pos_embed.reshape(arx[0] * arx[1], *_pos_embed.shape[2:])
+            x[idx, : arx[0] * arx[1]] += _pos_embed * self.gated_positional_embedding_gate.tanh()
+        return x
+
+
+def get_example_inputs():
+    x = torch.randn(1, MAX_NUM_CHUNKS, VISION_SEQ_LEN, VISION_HIDDEN_DIM, dtype=TORCH_DTYPE)
+    ar = torch.tensor([1, 1], dtype=torch.int32).view(1, 2)
+    return x, ar
+
+
+def test_apply_pos_embed():
+    setup_debug_env()
+
+    init_args = dict(
+        max_num_tiles=MAX_NUM_CHUNKS,
+        image_size=560,
+        patch_size=14,
+        width=VISION_HIDDEN_DIM,
+    )
+
+    cpu_model_meta = VisionEncoderMeta(**init_args)
+    save_checkpoint(cpu_model_meta)
+    cpu_model = VisionEncoderPosEmbedOnly(**init_args)
+    cpu_model.load_state_dict(load_checkpoint())
+
+    # Trace to get neuron model
+    example_inputs = get_example_inputs()
+    x, ar = example_inputs
+    neuron_model = trace_nxd_model(
+        VisionEncoderPosEmbedOnly, example_inputs, tp_degree=1, **init_args
+    )
+
+    # Test all possible aspect ratios (with max_num_chunks=4)
+    aspect_ratios = [[1, 1], [1, 2], [1, 3], [1, 4], [2, 1], [2, 2], [3, 1], [4, 1]]
+    for aspect_ratio in aspect_ratios:
+        print("Testing aspect ratio:", tuple(aspect_ratio))
+        ar = torch.tensor(aspect_ratio, dtype=torch.int32).view(1, 2)
+
+        # Compare Meta vs our implementation on CPU
+        x_out_meta = cpu_model_meta(x, ar)
+        x_out_cpu = cpu_model(x, ar)
+        assert torch.allclose(x_out_meta, x_out_cpu)
+        logger.info("Correctness test passing on CPU.")
+
+        x_out_xla = neuron_model(x, ar)
+        assert torch.allclose(x_out_meta, x_out_xla)
+        logger.info(
+            f"{x_out_meta.shape}, {x.sum()}, {x_out_meta.sum()}, {x_out_cpu.sum()}, {x_out_xla.sum()}"
+        )
+        logger.info("Correctness test passing on device.\n")
+
+    logger.info("ALL TESTS PASSING")
+
+
+if __name__ == "__main__":
+    test_apply_pos_embed()