princeton-nlp / llm-shearing Goto Github PK

LLM-Shearing/llmshearing/datasets/streaming_da │
│ taset.py:46 in generate_work                                                                     │
│                                                                                                  │
│    43 │   │   List[List[int]]: The epoch for each domain of data (num physical nodes,            │
│    44 │   │   ranks per node, workers per rank, batches per worker, batch size).                 │
│    45 │   """                                                                                    │
│ ❱  46 │   assert epoch == 0, "Currently only supports dynamic loading from each domain for onc   │
│    47 │   # Ensure that num_canonical_nodes has been set.                                        │
│    48 │   if dataset.num_canonical_nodes is None:                                                │
│    49 │   │   raise RuntimeError(f'`num_canonical_nodes` can never be None. ' +                  │
╰───────────────────────────────────────────
AssertionError: Currently only supports dynamic loading from each domain for once.

# Currently only supports dynamically loading data from each domain for once. 
        # Issues could occur if one domain of data is used up. 
        while True:
            proportion = self.proportion
            stream_id = np.random.choice(range(self.num_streams), 1, p=proportion)[0].item()
            domain_sample_id = sample_ids_per_stream[stream_id]
            domain_sample_id = domain_sample_id[self.used_num_samples_per_stream[stream_id] \
                                % self.samples_per_stream[stream_id]]
            self.used_num_samples_per_stream[stream_id] += 1
            yield self[domain_sample_id]
---
IndexError: index 24 is out of bounds for axis 0 with size 24

│   304 │   │   │   # Load the generated epoch shape from shared memory.                           │
│   305 │   │   │   name = _get_path(self._shm_prefix_int, EPOCH_SHAPE + f"_{stream_id}")          │
│   306 │   │   │   size = ndim * np.int64().nbytes                                                │
│ ❱ 307 │   │   │   shape_shm = SharedMemory(name=name, create=False, size=size, auto_cleanup=Fa   │
│   308 │   │   │   shape = tuple(np.ndarray(5, buffer=shape_shm.buf, dtype=np.int64))             │
│   309 │   │   │                                                                                  │
│   310 │   │   │   # Attach to the generated epoch data in shared memory.      
---
FileNotFoundError: [Errno 2] No such file or directory: '/000000_epoch_shape_0'
Exception ignored in atexit callback: <function Engine._close at 0x7fe6dc766a70>

def get_gpt_key_map_from_hf_to_composer(num_layers):
    """ get the keymap from hf to composer """
    key_map = {}
    key_map.update({"gpt_neox.embed_in.weight": "model.transformer.wte.weight",
                    "gpt_neox.final_layer_norm.weight": "model.transformer.ln_f.weight",
                    "gpt_neox.final_layer_norm.bias": "model.transformer.ln_f.bias",
                    "embed_out.weight": "model.transformer.output.weight",
                    "embed_out.bias": "model.transformer.output.bias"})
    for i in range(num_layers):
        key_map.update({
                        f"gpt_neox.layers.{i}.input_layernorm.weight": f"model.transformer.blocks.{i}.ln_1.weight",
                        f"gpt_neox.layers.{i}.post_attention_layernorm.weight": f"model.transformer.blocks.{i}.ln_2.weight",
                        f"gpt_neox.layers.{i}.attention.query_key_value.weight": f"model.transformer.blocks.{i}.attn.query_key_value.weight",
                        f"gpt_neox.layers.{i}.attention.dense.weight": f"model.transformer.blocks.{i}.attn.out_proj.weight",
                        f"gpt_neox.layers.{i}.mlp.dense_h_to_4h.weight": f"model.transformer.blocks.{i}.mlp.up_proj.weight",
                        f"gpt_neox.layers.{i}.mlp.dense_4h_to_h.weight": f"model.transformer.blocks.{i}.mlp.down_proj.weight",
                        f"gpt_neox.layers.{i}.input_layernorm.bias": f"model.transformer.blocks.{i}.ln_1.bias",
                        f"gpt_neox.layers.{i}.post_attention_layernorm.bias": f"model.transformer.blocks.{i}.ln_2.bias",
                        f"gpt_neox.layers.{i}.attention.query_key_value.bias": f"model.transformer.blocks.{i}.attn.query_key_value.bias",
                        f"gpt_neox.layers.{i}.attention.dense.bias": f"model.transformer.blocks.{i}.attn.out_proj.bias",
                        f"gpt_neox.layers.{i}.mlp.dense_h_to_4h.bias": f"model.transformer.blocks.{i}.mlp.up_proj.bias",
                        f"gpt_neox.layers.{i}.mlp.dense_4h_to_h.bias": f"model.transformer.blocks.{i}.mlp.down_proj.bias",
                        f"gpt_neox.layers.{i}.attention.rotary_emb.inv_freq": f"model.transformer.blocks.{i}.attn.rotary_emb.inv_freq",
                       })
    return key_map

def construct_example_cfg(model_size, path=None, add_l0_module=False):
    """ construct example cfg for mosaicml llama models """
    if model_size == "1.4b":
        cfg = om.create({"name": "mosaic_pythia_1", "init_device": "cpu", "d_model": 2048, 
                         "n_heads": 16, "n_layers": 24, "intermediate_size": 8192,
                         "rotary_pct": 0.25, "rotary_emb_base": 10000
                         })

    # add default values
    cfg = om.merge(cfg, om.create({"max_seq_len": 2048, "vocab_size": 50304, "init_std": 0.02, "attn_pdrop": 0.0, "resid_pdrop": 0.0, "emb_pdrop": 0.0, "attn_impl": "norm", "layer_norm_eps": 1e-5}))
    if add_l0_module:
        cfg["l0_module"] = {"start_sparsity": 0, "target_sparsity": 0.6, "pruning_modules": ["head", "head_layer", "mlp", "intermediate", "hidden"], "lagrangian_warmup_steps": "320ba"}
    return cfg

(Pdb) composer_model
ComposerMosaicPythia(
  (model): PythiaModel(
    (transformer): ModuleDict(
      (wte): PythiaEmbedding(50304, 2048)
      (blocks): ModuleList(
        (0-23): 24 x PythiaBlock(
          (ln_1): CoFiLayerNorm((2048,), eps=1e-05, elementwise_affine=True)
          (attn): PythiaAttention(
            (query_key_value): Linear(in_features=2048, out_features=6144, bias=True)
            (out_proj): Linear(in_features=2048, out_features=2048, bias=True)
            (rotary_emb): RotaryEmbedding()
          )
          (ln_2): CoFiLayerNorm((2048,), eps=1e-05, elementwise_affine=True)
          (mlp): PythiaMLP(
            (down_proj): Linear(in_features=8192, out_features=2048, bias=True)
            (up_proj): Linear(in_features=2048, out_features=8192, bias=True)
          )
        )
      )
      (output): Linear(in_features=2048, out_features=50304, bias=False)
      (ln_f): CoFiLayerNorm((2048,), eps=1e-05, elementwise_affine=True)
    )
  )
)
(Pdb) hf_model
GPTNeoXForCausalLM(
  (gpt_neox): GPTNeoXModel(
    (embed_in): Embedding(50304, 2048)
    (emb_dropout): Dropout(p=0.0, inplace=False)
    (layers): ModuleList(
      (0-23): 24 x GPTNeoXLayer(
        (input_layernorm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)
        (post_attention_layernorm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)
        (post_attention_dropout): Dropout(p=0.0, inplace=False)
        (post_mlp_dropout): Dropout(p=0.0, inplace=False)
        (attention): GPTNeoXAttention(
          (rotary_emb): GPTNeoXRotaryEmbedding()
          (query_key_value): Linear(in_features=2048, out_features=6144, bias=True)
          (dense): Linear(in_features=2048, out_features=2048, bias=True)
          (attention_dropout): Dropout(p=0.0, inplace=False)
        )
        (mlp): GPTNeoXMLP(
          (dense_h_to_4h): Linear(in_features=2048, out_features=8192, bias=True)
          (dense_4h_to_h): Linear(in_features=8192, out_features=2048, bias=True)
          (act): GELUActivation()
        )
      )
    )
    (final_layer_norm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)
  )
  (embed_out): Linear(in_features=2048, out_features=50304, bias=False)
)

> /root/personal-code/LLM-Shearing/llmshearing/data/sample.py(69)<module>()
-> folder_eval_target = args.eval_seq
(Pdb) c
Traceback (most recent call last):
  File "/root/personal-code/LLM-Shearing/llmshearing/data/sample.py", line 70, in <module>
    num_sample_each_file = max(1, folder_eval_target // len(selected) + 1)
ZeroDivisionError: integer division or modulo by zero```

Building train loader...
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│                          /LLM-Shearing/llmshearing/train.py:317 in <module>                      │
│                                                                                                  │
│   314 │   os.makedirs(save_dir, exist_ok=True)                                                   │
│   315 │   torch.save(cfg, save_dir + "/config.pt")                                               │
│   316 │                                                                                          │
│ ❱ 317 │   main(cfg)                                                                              │
│   318                                                                                            │
│   319                                                                                            │
│                                                                                                  │
│                          /LLM-Shearing/llmshearing/train.py:201 in main                          │
│                                                                                                  │
│   198 │                                                                                          │
│   199 │   # Dataloaders                                                                          │
│   200 │   print('Building train loader...')                                                      │
│ ❱ 201 │   train_loader = build_text_dataloader(cfg.train_loader,                                 │
│   202 │   │   │   │   │   │   │   │   │   │    cfg.device_train_batch_size,                      │
│   203 │   │   │   │   │   │   │   │   │   │    cfg.callbacks.data_loading.dynamic,               │
│   204 │   │   │   │   │   │   │   │   │   │    cfg.callbacks.data_loading.set_names,             │
│                                                                                                  │
│                          /LLM-Shearing/llmshearing/datasets/load_text_dataloader.py:36 in        │
│ build_text_dataloader                                                                            │
│                                                                                                  │
│    33 │   """                                                                                    │
│    34 │                                                                                          │
│    35 │   if dynamic:                                                                            │
│ ❱  36 │   │   dataset = TextDynamicStreamingDataset(local=cfg.dataset.local,                     │
│    37 │   │   │   │   │   │   │   │   │   │   │     max_seq_len=cfg.dataset.max_seq_len,         │
│    38 │   │   │   │   │   │   │   │   │   │   │     batch_size=device_batch_size,                │
│    39 │   │   │   │   │   │   │   │   │   │   │     shuffle=cfg.dataset.get(                     │
│                                                                                                  │
│                          /LLM-Shearing/llmshearing/datasets/streaming_dataset.py:415 in __init__ │
│                                                                                                  │
│   412 │   │   │   │    is_uint16: bool = False):                                                 │
│   413 │   │                                                                                      │
│   414 │   │   # Build Dataset                                                                    │
│ ❱ 415 │   │   super().__init__(local=local,                                                      │
│   416 │   │   │   │   │   │    shuffle=shuffle,                                                  │
│   417 │   │   │   │   │   │    shuffle_seed=shuffle_seed,                                        │
│   418 │   │   │   │   │   │    num_canonical_nodes=num_canonical_nodes,                          │
│                                                                                                  │
│                          /LLM-Shearing/llmshearing/datasets/streaming_dataset.py:114 in __init__ │
│                                                                                                  │
│   111 │   │   │   │    proportion: List[float] = None) -> None:                                  │
│   112 │   │                                                                                      │
│   113 │   │   streams = [Stream(local=local, split=set_name, repeat=1.0) for set_name in set_n   │
│ ❱ 114 │   │   super().__init__(streams=streams,                                                  │
│   115 │   │   │   │   │   │    split=None,                                                       │
│   116 │   │   │   │   │   │    num_canonical_nodes=num_canonical_nodes,                          │
│   117 │   │   │   │   │   │    batch_size=batch_size,                                            │
│                                                                                                  │
│                          /lib/python3.10/site-packages/streaming/base/dataset.py:443 in    │
│ __init__                                                                                         │
│                                                                                                  │
│    440 │   │   self.sample_offset_per_stream = np.zeros(self.num_streams, np.int64)              │
│    441 │   │   self.samples_per_stream = np.zeros(self.num_streams, np.int64)                    │
│    442 │   │   for stream_id, stream in enumerate(self.streams):                                 │
│ ❱  443 │   │   │   stream_shards = stream.get_shards(world)                                      │
│    444 │   │   │   num_stream_samples = sum(map(len, stream_shards))                             │
│    445 │   │   │   if not num_stream_samples:                                                    │
│    446 │   │   │   │   index_filename = os.path.join(stream.local, stream.split, get_index_base  │
│                                                                                                  │
│                                /lib/python3.10/site-packages/streaming/base/stream.py:437 in     │
│ get_shards                                                                                       │
│                                                                                                  │
│   434 │   │   │   │   │   os.rename(tmp_filename, filename)                                      │
│   435 │   │   │   │   else:                                                                      │
│   436 │   │   │   │   │   if not os.path.exists(filename):                                       │
│ ❱ 437 │   │   │   │   │   │   raise RuntimeError(f'No `remote` provided, but local file {filen   │
│   438 │   │   │   │   │   │   │   │   │   │      'does not exist either')                        │
│   439 │   │   │   else:                                                                          │
│   440 │   │   │   │   wait_for_file_to_exist(                                                    │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: No `remote` provided, but local file /for_prune/cc/index.json does not exist either

[metric][batch=0]: time/epoch: 2365 
[metric][batch=0]: metrics/train/LanguageCrossEntropy: nan 
[metric][batch=0]: metrics/train/Perplexity: nan 
[metric][batch=0]: metrics/train/ArXiv_LanguageCrossEntropy: nan 
[metric][batch=0]: metrics/train/ArXiv_count: 0 
[metric][batch=0]: metrics/train/Books_LanguageCrossEntropy: nan 
[metric][batch=0]: metrics/train/Books_count: 0 
[metric][batch=0]: metrics/train/Wikipedia_LanguageCrossEntropy: nan 
[metric][batch=0]: metrics/train/Wikipedia_count: 0 
[metric][batch=0]: time/epoch: 2366 
[metric][batch=0]: metrics/train/LanguageCrossEntropy: nan 
[metric][batch=0]: metrics/train/Perplexity: nan 
[metric][batch=0]: metrics/train/ArXiv_LanguageCrossEntropy: nan 
[metric][batch=0]: metrics/train/ArXiv_count: 0 
[metric][batch=0]: metrics/train/Books_LanguageCrossEntropy: nan 
[metric][batch=0]: metrics/train/Books_count: 0 
[metric][batch=0]: metrics/train/Wikipedia_LanguageCrossEntropy: nan 
[metric][batch=0]: metrics/train/Wikipedia_count: 0 
.... repeat

# learning setup
lr=1e-4 # learning rate for the main parameters
max_duration=3200ba # 0.42B tokens
save_interval=3200ba # save in the end
t_warmup=320ba # 10% learning rate warmup 

# dynamic loading setup
dynamic=True
# set_names=[cc,github,book,stackexchange,wiki,arxiv,c4-rp] # domain names
set_names=[ArXiv,Books,Wikipedia] # domain names
# proportion=[0.67,0.045,0.045,0.02,0.045,0.025,0.15] # initial proportion of RP, make sure that the sum(proportion) = 1
proportion=[0.4,0.3,0.3] 
# doremi: update weights with exponential descent
# constant: keep the weights constant
update_type=doremi 
if [[ $to_model == 1.3b ]]; then
    # target_loss=[1.9643,0.7459,2.1393,1.6117,1.7590,1.4449,2.1251] # 1.3b predicted loss from scaling law
    target_loss=[1.4449,2.1393,1.7590]
else
    # target_loss=[1.8712,0.6883,2.0325,1.5353,1.6297,1.3560,2.0328] # 2.7b predicted loss from scaling law
    target_loss=[1.3560,2.0325,1.6297]
fi
eval_split_name=eval_merge # eval on all domains
eval_target_model=false # evaluate on the current model, not the target model, otherwise the loss will be inaccurate
eval_interval=50ba # eval every 50 batches and update the loading proportion


# pruning setup
lag_lr=1.0 # learning rate or l0_module
lagr_warmup=640ba # 20% sparsity warmup
if [[ $to_model == 1.3b ]]; then
    target_d_model=2048; target_n_heads=16; target_n_layers=24; target_intermediate_size=5504
elif [[ $to_model == 3b ]]; then
    target_d_model=2560; target_n_heads=20; target_n_layers=32; target_intermediate_size=6912
fi

composer $TRAIN_SCRIPT \
    $config_file \
    run_name=${run_name} \
    data_local=${data_local} \
    eval_loader.dataset.split=${eval_split_name} \
    global_train_batch_size=${global_train_batch_size} \
    device_train_microbatch_size=${device_train_microbatch_size} \
    device_eval_batch_size=${device_eval_batch_size} \
    max_seq_len=${max_seq_len} \
    max_duration=${max_duration} \
    eval_first=false \
    scheduler.t_warmup=${t_warmup} \
    save_folder=${save_dir} \
    loggers.wandb.init_kwargs.dir=${wandb_dir} \
    eval_interval=${eval_interval} \
    save_interval=${save_interval} \
    optimizer.lr=${lr} \
    optimizer.lag_lr=${lag_lr} \
    model.l0_module.lagrangian_warmup_steps=${lagr_warmup} \
    model.l0_module.pruning_modules='[head,intermediate,layer,hidden]' \
    model.l0_module.eval_target_model=${eval_target_model} \
    model.l0_module.target_model.d_model=${target_d_model} \
    model.l0_module.target_model.n_heads=${target_n_heads} \
    model.l0_module.target_model.n_layers=${target_n_layers} \
    model.l0_module.target_model.intermediate_size=${target_intermediate_size} \
    callbacks.data_loading.dynamic=${dynamic} \
    callbacks.data_loading.set_names=${set_names} \
    callbacks.data_loading.proportion=${proportion} \
    callbacks.data_loading.update_type=${update_type} \
    callbacks.data_loading.target_loss=${target_loss} \
    train_loader.num_workers=0 \
    train_loader.prefetch_factor=null \
    train_loader.persistent_workers=false \
    autoresume=false

 def forward(self, batch):
        input_ids = batch["input_ids"]
        key_padding_mask = batch["attention_mask"].bool() if "attention_mask" in batch else None
        pruned_steps = batch.get("pruned_steps", None)
        if pruned_steps is not None:
            pruned_steps = pruned_steps[0].item()
        zs = {key: batch[key] for key in batch if "_z" in key} # why those parameters coming from batch?
        model_output = self.model(
            input_ids=input_ids, key_padding_mask=key_padding_mask, pruned_steps=pruned_steps, **zs
        )
        return model_output

# Specify $PROJ_DIR in scripts/launch.sh and scripts/srun_launch.sh if using slurm

test=True

from_model=7b # source model size
to_model=3b # target model size
config_file=${PROJ_DIR}/llmshearing/configs/llama2/${from_model}.yaml
path=$MODEL_PATH/llama-2-7b-composer/state_dict.pt

# data setup
data_local=${DATA_DIR}

# basic setup
max_seq_len=4096
device_train_microbatch_size=4
global_train_batch_size=32
device_eval_batch_size=8

# learning setup
lr=1e-4 # learning rate for the main parameters
# max_duration=3200ba # 0.42B tokens
# save_interval=3200ba # save in the end
# t_warmup=320ba # 10% learning rate warmup 

max_duration=3000ba # 0.39B tokens
save_interval=3000ba # save in the end
t_warmup=300ba # 10% learning rate warmup 

# dynamic loading setup
dynamic=True
set_names=[cc,github,book,stackexchange,wiki,arxiv,c4-rp] # domain names
proportion=[0.67,0.045,0.045,0.02,0.045,0.025,0.15] # initial proportion of RP, make sure that the sum(proportion) = 1
# doremi: update weights with exponential descent
# constant: keep the weights constant
update_type=doremi 
if [[ $to_model == 1.3b ]]; then
    target_loss=[1.9643,0.7459,2.1393,1.6117,1.7590,1.4449,2.1251] # 1.3b predicted loss from scaling law
else
    target_loss=[1.8712,0.6883,2.0325,1.5353,1.6297,1.3560,2.0328] # 2.7b predicted loss from scaling law
fi
eval_split_name=../eval/eval_merge # eval on all domains
eval_target_model=false # evaluate on the current model, not the target model, otherwise the loss will be inaccurate
eval_interval=50ba # eval every 50 batches and update the loading proportion


# pruning setup
lag_lr=1.0 # learning rate or l0_module
lagr_warmup=640ba # 20% sparsity warmup
if [[ $to_model == 1.3b ]]; then
    target_d_model=2048; target_n_heads=16; target_n_layers=24; target_intermediate_size=5504
elif [[ $to_model == 3b ]]; then
    target_d_model=2560; target_n_heads=20; target_n_layers=32; target_intermediate_size=6912
fi

# save directroy
run_name=llama2_${from_model}_pruning_scaling_${update_type}_to${to_model}_sl${max_seq_len}
save_dir=${OUTPUT_DIR}/${run_name}
wandb_dir=${save_dir} # save locally

if [[ $test == True ]]; then t=00-01:00:00; else t=01-00:00:00; fi

# Run in bash, it will automatically use resources available in the current environment
# composer $TRAIN_SCRIPT \

# Run with slurm    
# sbatch -p cli \
#     --job-name ${run_name} \
#     --nodes=4 \
#     --gpus-per-node=2 \
#     --mem=512gb \
#     --cpus-per-task=8 \
#     --time $t \
composer $TRAIN_SCRIPT \
    $config_file \
    run_name=${run_name} \
    data_local=${data_local} \
    eval_loader.dataset.split=${eval_split_name} \
    global_train_batch_size=${global_train_batch_size} \
    device_train_microbatch_size=${device_train_microbatch_size} \
    device_eval_batch_size=${device_eval_batch_size} \
    max_seq_len=${max_seq_len} \
    max_duration=${max_duration} \
    eval_first=false \
    scheduler.t_warmup=${t_warmup} \
    save_folder=${save_dir} \
    loggers.wandb.init_kwargs.dir=${wandb_dir} \
    eval_interval=${eval_interval} \
    save_interval=${save_interval} \
    optimizer.lr=${lr} \
    optimizer.lag_lr=${lag_lr} \
    model.path=${path} \
    model.l0_module.lagrangian_warmup_steps=${lagr_warmup} \
    model.l0_module.pruning_modules='[head,intermediate,layer,hidden]' \
    model.l0_module.eval_target_model=${eval_target_model} \
    model.l0_module.target_model.d_model=${target_d_model} \
    model.l0_module.target_model.n_heads=${target_n_heads} \
    model.l0_module.target_model.n_layers=${target_n_layers} \
    model.l0_module.target_model.intermediate_size=${target_intermediate_size} \
    callbacks.data_loading.dynamic=${dynamic} \
    callbacks.data_loading.set_names=${set_names} \
    callbacks.data_loading.proportion=${proportion} \
    callbacks.data_loading.update_type=${update_type} \
    callbacks.data_loading.target_loss=${target_loss} \
    train_loader.num_workers=0 \
    train_loader.prefetch_factor=null \
    train_loader.persistent_workers=false \
    autoresume=true

User
log(event, f'Running callback {type(cb).__na │
│ ❱ 468 │   │   │   │   cb.run_event(event, self.state, self.logger)             │
│   469 │                                                                        │
│   470 │   def _run_loggers(self, event: Union[Event, str]):                    │
│   471 │   │   loggers = [callback for callback in self.state.callbacks if isin │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/composer/core/call │
│ back.py:96 in run_event                                                        │
│                                                                                │
│    93 │   │   │   logger (Logger): The logger.                                 │
│    94 │   │   """                                                              │
│    95 │   │   event_cb = getattr(self, event.value)                            │
│ ❱  96 │   │   return event_cb(state, logger)                                   │
│    97 │                                                                        │
│    98 │   def init(self, state: State, logger: Logger) -> None:                │
│    99 │   │   """Called on the :attr:`.Event.INIT` event.                      │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/composer/callbacks │
│ /checkpoint_saver.py:294 in batch_checkpoint                                   │
│                                                                                │
│   291 │   def batch_checkpoint(self, state: State, logger: Logger):            │
│   292 │   │   assert callable(self.save_interval)                              │
│   293 │   │   if self.save_interval(state, Event.BATCH_CHECKPOINT) and self.la │
│ ❱ 294 │   │   │   self._save_checkpoint(                                       │
│   295 │   │   │   │   state,                                                   │
│   296 │   │   │   │   logger,                                                  │
│   297 │   │   │   )                                                            │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/composer/callbacks │
│ /checkpoint_saver.py:332 in _save_checkpoint                                   │
│                                                                                │
│   329 │   │   # save the checkpoint to the filename                            │
│   330 │   │   filename_with_placeholders = self.filename.format(state, is_deep │
│   331 │   │                                                                    │
│ ❱ 332 │   │   saved_path = checkpoint.save_checkpoint(                         │
│   333 │   │   │   state=state,                                                 │
│   334 │   │   │   filename=filename_with_placeholders,                         │
│   335 │   │   │   weights_only=self.weights_only,                              │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/composer/utils/che │
│ ckpoint.py:761 in save_checkpoint                                              │
│                                                                                │
│   758 │   │   }                                                                │
│   759 │   else:                                                                │
│   760 │   │   state_dict = {                                                   │
│ ❱ 761 │   │   │   'state': state.state_dict(),                                 │
│   762 │   │   │   'rng': reproducibility.get_rng_state(),                      │
│   763 │   │   }                                                                │
│   764                                                                          │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/composer/core/stat │
│ e.py:891 in state_dict                                                         │
│                                                                                │
│    888 │   │   │   if attribute_name == 'dataset_state':                       │
│    889 │   │   │   │   serialized_value = self._dataset_state_dict()           │
│    890 │   │   │   elif attribute_name == 'model':                             │
│ ❱  891 │   │   │   │   serialized_value = self.get_model_state_dict()          │
│    892 │   │   │   elif attribute_name == 'optimizers':                        │
│    893 │   │   │   │   optimizer = ensure_tuple(attribute_value)[              │
│    894 │   │   │   │   │   0]  # Let's stop pretending. We don't support more  │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/composer/core/stat │
│ e.py:868 in get_model_state_dict                                               │
│                                                                                │
│    865 │   │   """                                                             │
│    866 │   │   if self.fsdp_enabled and self.fsdp_state_dict_type is not None: │
│    867 │   │   │   with fsdp_state_dict_type_context(self.model, state_dict_ty │
│ ❱  868 │   │   │   │   model_state_dict = self.model.state_dict()              │
│    869 │   │   else:                                                           │
│    870 │   │   │   model_state_dict = self.model.state_dict()                  │
│    871                                                                         │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/nn/modules/m │
│ odule.py:1818 in state_dict                                                    │
│                                                                                │
│   1815 │   │   self._save_to_state_dict(destination, prefix, keep_vars)        │
│   1816 │   │   for name, module in self._modules.items():                      │
│   1817 │   │   │   if module is not None:                                      │
│ ❱ 1818 │   │   │   │   module.state_dict(destination=destination, prefix=prefi │
│   1819 │   │   for hook in self._state_dict_hooks.values():                    │
│   1820 │   │   │   hook_result = hook(self, destination, prefix, local_metadat │
│   1821 │   │   │   if hook_result is not None:                                 │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/nn/modules/m │
│ odule.py:1815 in state_dict                                                    │
│                                                                                │
│   1812 │   │   if hasattr(destination, "_metadata"):                           │
│   1813 │   │   │   destination._metadata[prefix[:-1]] = local_metadata         │
│   1814 │   │                                                                   │
│ ❱ 1815 │   │   self._save_to_state_dict(destination, prefix, keep_vars)        │
│   1816 │   │   for name, module in self._modules.items():                      │
│   1817 │   │   │   if module is not None:                                      │
│   1818 │   │   │   │   module.state_dict(destination=destination, prefix=prefi │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/nn/modules/m │
│ odule.py:1722 in _save_to_state_dict                                           │
│                                                                                │
│   1719 │   │   │   │   module                                                  │
│   1720 │   │   """                                                             │
│   1721 │   │   for hook in self._state_dict_pre_hooks.values():                │
│ ❱ 1722 │   │   │   hook(self, prefix, keep_vars)                               │
│   1723 │   │                                                                   │
│   1724 │   │   for name, param in self._parameters.items():                    │
│   1725 │   │   │   if param is not None:                                       │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/utils/_conte │
│ xtlib.py:115 in decorate_context                                               │
│                                                                                │
│   112 │   @functools.wraps(func)                                               │
│   113 │   def decorate_context(*args, **kwargs):                               │
│   114 │   │   with ctx_factory():                                              │
│ ❱ 115 │   │   │   return func(*args, **kwargs)                                 │
│   116 │                                                                        │
│   117 │   return decorate_context                                              │
│   118                                                                          │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/distributed/ │
│ fsdp/_state_dict_utils.py:669 in _pre_state_dict_hook                          │
│                                                                                │
│   666 │   │   StateDictType.LOCAL_STATE_DICT: _local_pre_state_dict_hook,      │
│   667 │   │   StateDictType.SHARDED_STATE_DICT: _sharded_pre_state_dict_hook,  │
│   668 │   }                                                                    │
│ ❱ 669 │   _pre_state_dict_hook_fn[fsdp_state._state_dict_type](                │
│   670 │   │   fsdp_state,                                                      │
│   671 │   │   module,                                                          │
│   672 │   │   *args,                                                           │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/distributed/ │
│ fsdp/_state_dict_utils.py:271 in _full_pre_state_dict_hook                     │
│                                                                                │
│   268 │   in ``nn.Module``.                                                    │
│   269 │   """                                                                  │
│   270 │   _common_pre_state_dict_hook(module, fsdp_state)                      │
│ ❱ 271 │   _common_unshard_pre_state_dict_hook(                                 │
│   272 │   │   module,                                                          │
│   273 │   │   fsdp_state,                                                      │
│   274 │   │   offload_to_cpu=fsdp_state._state_dict_config.offload_to_cpu,     │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/distributed/ │
│ fsdp/_state_dict_utils.py:143 in _common_unshard_pre_state_dict_hook           │
│                                                                                │
│   140 │   Performs the pre-state_dict tasks shared by all state_dict types tha │
│   141 │   ``_unshard_fsdp_state_params()``. FULL_STATE_DICT and SHARDED_STATE_ │
│   142 │   """                                                                  │
│ ❱ 143 │   _enter_unshard_params_ctx(                                           │
│   144 │   │   module,                                                          │
│   145 │   │   fsdp_state,                                                      │
│   146 │   │   writeback=False,                                                 │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/distributed/ │
│ fsdp/_state_dict_utils.py:109 in _enter_unshard_params_ctx                     │
│                                                                                │
│   106 │   │   offload_to_cpu=offload_to_cpu,                                   │
│   107 │   │   with_grads=with_grads,                                           │
│   108 │   )                                                                    │
│ ❱ 109 │   fsdp_state._unshard_params_ctx[module].__enter__()                   │
│   110                                                                          │
│   111                                                                          │
│   112 @no_type_check                                                           │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/contextlib.py:135 in __enter__   │
│                                                                                │
│   132 │   │   # they are only needed for recreation, which is not possible any │
│   133 │   │   del self.args, self.kwds, self.func                              │
│   134 │   │   try:                                                             │
│ ❱ 135 │   │   │   return next(self.gen)                                        │
│   136 │   │   except StopIteration:                                            │
│   137 │   │   │   raise RuntimeError("generator didn't yield") from None       │
│   138                                                                          │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/distributed/ │
│ fsdp/_unshard_param_utils.py:171 in _unshard_fsdp_state_params                 │
│                                                                                │
│   168 │   This unshards the parameters for a single FSDP state ``state`` that  │
│   169 │   corresponds to ``module``.                                           │
│   170 │   """                                                                  │
│ ❱ 171 │   _validate_unshard_params_args(                                       │
│   172 │   │   state, writeback, rank0_only, offload_to_cpu, with_grads         │
│   173 │   )                                                                    │
│   174 │   torch.cuda.synchronize()                                             │
│                                                                                │
│ /root/miniconda3/envs/shearing/lib/python3.10/site-packages/torch/distributed/ │
│ fsdp/_unshard_param_utils.py:140 in _validate_unshard_params_args              │
│                                                                                │
│   137 │   if offload_to_cpu and any(                                           │
│   138 │   │   not handle.uses_sharded_strategy for handle in state._handles    │
│   139 │   ):                                                                   │
│ ❱ 140 │   │   raise NotImplementedError(                                       │
│   141 │   │   │   "offload_to_cpu=True and NO_SHARD is not supported yet"      │
│   142 │   │   )                                                                │
│   143 │   if writeback and rank0_only:                                         │
╰────────────────────────────────────────────────────────────────────────────────╯
NotImplementedError: offload_to_cpu=True and NO_SHARD is not supported yet

MODEL_PATH='models/Llama-2-7b-composer/pruned-state_dict.pt'
OUTPUT_PATH='models/Llama-2-7b-composer/hf-state_dict'
MODEL_CLASS='LlamaForCausalLM'
HIDDEN_SIZE=2048
NUM_ATTENTION_HEADS=16
NUM_HIDDEN_LAYERS=24
INTERMEDIATE_SIZE=5504
MODEL_NAME='Sheared-Llama-1.3B'

!python3 -m llmshearing.utils.composer_to_hf $MODEL_PATH $OUTPUT_PATH \
model_class=$MODEL_CLASS \
hidden_size=$HIDDEN_SIZE \
num_attention_heads=$NUM_ATTENTION_HEADS \
num_hidden_layers=$NUM_HIDDEN_LAYERS \
intermediate_size=$INTERMEDIATE_SIZE \
num_key_value_heads=$NUM_ATTENTION_HEADS \
_name_or_path=$MODEL_NAME

│ /tf/LLM-Shearing/llmshearing/utils/composer_to_hf.py:108 in <module>         │
│                                                                              │
│   105 if __name__ == "__main__":                                             │
│   106 │   composer_model_path, output_path, other_args = sys.argv[1], sys.ar │
│   107 │   cli_cfg = om.from_cli(other_args)                                  │
│ ❱ 108 │   save_composer_to_hf(composer_model_path, output_path, cli_cfg)     │
│   109 │   #save_hf_to_composer(composer_model_path, output_path)             │
│   110                                                                        │
│                                                                              │
│ /tf/LLM-Shearing/llmshearing/utils/composer_to_hf.py:90 in                   │
│ save_composer_to_hf                                                          │
│                                                                              │
│    87 def save_composer_to_hf(composer_model_path, output_path=None, model_c │
│    88 │   """ convert composer ckpt's weights to huggingface """             │
│    89 │                                                                      │
│ ❱  90 │   weights = torch.load(composer_model_path)["state"]["model"]        │
│    91 │   num_layers = get_layer_num_from_weights(weights)                   │
│    92 │   keymap = get_key_map_from_composer_to_hf(num_layers)               │
│    93 │   hf_weights = {keymap[key]: weights[key] for key in weights if "rot │
╰──────────────────────────────────────────────────────────────────────────────╯
KeyError: 'state'

def param_init_fn(self, module):
        """ Initialize the parameters for masking variables. """
        mean = math.log(1 - self.droprate_init) - math.log(self.droprate_init)
        mean = 5
        if isinstance(module, nn.Parameter):
            module.data.normal_(mean, 1e-2)
        else:
            for tensor in module.parameters():
                tensor.data.normal_(mean, 1e-2)