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)

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

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

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

# Install fa4 with `uv add git+https://github.com/Dao-AILab/flash-attention.git@main#subdirectory=flash_attn/cute`

import torch
from torch.library import Library

from flash_attn.cute.interface import _flash_attn_fwd, _flash_attn_bwd

# I guess the garbage collector removes the library object so it unregisters the implementation if not kept alive?
_lib = None

S1ro1

---
# Source: pd-stack/templates/router.yaml
apiVersion: v1
kind: Service
metadata:
  name: glm47-router
  namespace: dev-matej
  labels:
    app: glm47
    chart: pd-stack