S1ro1

import torch
from torch import nn
from torch.distributed.tensor.placement_types import Replicate, Shard
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.parallel import parallelize_module


def dist_print(*args, **kwargs):

S1ro1

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
enable_cpu_affinity: false
fsdp_config:
  fsdp_activation_checkpointing: false
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false

S1ro1

from transformers import AutoModelForCausalLM
from accelerate import Accelerator
import torch

torch.cuda.memory._record_memory_history()

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
accelerator = Accelerator()

model = AutoModelForCausalLM.from_pretrained(model_id)

S1ro1

import torch


def print_test_end():
    print("---------------")


def test_vectors_bwd():
    print("TEST VECTORS BTW")
    a = torch.tensor([[1.0, -2.0, 3.0]], requires_grad=True)