O repositório do TrafficLLM , uma estrutura de adaptação universal de LLM para aprender a representação robusta de tráfego para todos os LLM de código aberto em cenários do mundo real e aprimorar a generalização em diversas tarefas de análise de tráfego.
Nota: este código é baseado em ChatGLM2 e Llama2. Muito obrigado aos autores.
[2024.10.28] Atualizamos o código de adaptação para usar GLM4 para construir TrafficLLM, que tem ajuste e velocidade de inferência mais rápidos do que ChatGLM2. Vá para Adapt2GLM4 para mais detalhes.
TrafficLLM é construído sobre uma estrutura sofisticada de ajuste fino usando linguagem natural e dados de tráfego, que propõe as seguintes técnicas para aumentar a utilidade de grandes modelos de linguagem na análise de tráfego de rede.
Tokenização de domínio de tráfego. Para superar a lacuna de modalidade entre a linguagem natural e os dados de tráfego heterogêneos, o TrafficLLM introduz a tokenização de domínio de tráfego para processar as diversas entradas de tarefas de detecção e geração de tráfego para adaptação do LLM. Este mecanismo estende efetivamente o tokenizer nativo do LLM, treinando o modelo de tokenização especializado em corpora de domínio de tráfego em grande escala.
Pipeline de ajuste de dois estágios. TrafficLLM emprega um pipeline de ajuste de dois estágios para obter o aprendizado robusto de representação do LLM em diferentes tarefas de domínio de tráfego. O pipeline treina o LLM para entender instruções e aprender padrões de tráfego relacionados a tarefas em diferentes estágios, o que se baseia na compreensão de tarefas do TrafficLLM e nas habilidades de raciocínio de tráfego para diversas tarefas de detecção e geração de tráfego.
Adaptação extensível com ajuste fino de parâmetro eficaz (EA-PEFT). Para adaptar o LLM para generalização a novos ambientes de tráfego, o TrafficLLM propõe uma adaptação extensível com ajuste fino de parâmetro efetivo (EA-PEFT) para atualizar parâmetros do modelo com baixa sobrecarga. A técnica divide as capacidades do modelo em diferentes modelos PEFT, o que ajuda a minimizar os custos em cenários dinâmicos gerados por mudanças no padrão de tráfego.
Lançamos os conjuntos de dados de treinamento do TrafficLLM, que contêm mais de 0,4 milhão de dados de tráfego e 9 mil instruções humanas para adaptação do LLM em diferentes tarefas de análise de tráfego.
Instruction Datasets : os conjuntos de dados de instrução são usados para ajudar o LLM a aprender o conhecimento do domínio de tarefas de detecção ou geração de tráfego e entender qual tarefa deve ser conduzida em diferentes cenários.
Traffic Datasets : os conjuntos de dados de tráfego contêm os dados de ajuste de tráfego que extraímos dos conjuntos de dados de tráfego públicos, o que ajuda o LLM a aprender o padrão de tráfego em diferentes tarefas downstream.
Para construir o corpus de linguagem natural como instruções humanas no TrafficLLM, coletamos 9.209 instruções específicas de tarefas supervisionadas por especialistas e assistentes de IA. As estatísticas são mostradas a seguir:
| Tarefas Principais | Tarefas posteriores | Abrev. | #Amostra |
|---|---|---|---|
| Detecção de tráfego | Detecção de tráfego de malware | MTD | 1,0K |
| Detecção de botnets | BND | 1,1K | |
| Detecção de DoH malicioso | MDD | 0,6K | |
| Detecção de ataques na Web | WAD | 0,6K | |
| Detecção de ataque APT | DAA | 0,6K | |
| Detecção de VPN criptografada | DVE | 1,2K | |
| Detecção de comportamento Tor | A definir | 0,6K | |
| Classificação de aplicativos criptografados | EAC | 0,6K | |
| Impressão digital do site | WF | 0,6K | |
| Deriva de Conceito | CD | 0,6K | |
| Geração de tráfego | Geração de tráfego de malware | MTG | 0,6K |
| Geração de tráfego de botnet | BTG | 0,1K | |
| Geração de VPN criptografada | EVG | 0,4K | |
| Geração de aplicativos criptografados | EAG | 0,6K |
Para avaliar o desempenho do TrafficLLM em vários cenários de rede, extraímos mais de 0,4 milhão de dados de ajuste de conjuntos de dados de tráfego disponíveis ao público para medir a capacidade do TrafficLLM de detectar ou gerar tráfego malicioso e benigno. As estatísticas são mostradas a seguir:
| Conjuntos de dados | Tarefas | Abrev. | #Amostra |
|---|---|---|---|
| USTCTFC 2016 | Detecção de tráfego de malware | MTD | 50,7 mil |
| Botnet ISCX 2014 | Detecção de botnets | BND | 25,0 mil |
| DoHBrw 2020 | Detecção de DoH malicioso | MDD | 47,8K |
| CSIC 2010 | Detecção de ataques na Web | WAD | 34,5 mil |
| DAPT 2020 | Detecção de ataque APT | DAA | 10,0K |
| VPN ISCX 2016 | Detecção de VPN criptografada | DVE | 64,8K |
| ISCXTor 2016 | Detecção de comportamento Tor | A definir | 40,0K |
| CSTNET 2023 | Classificação de aplicativos criptografados | EAC | 97,6K |
| CW-100 2018 | Impressão digital do site | WF | 7,4K |
| APP-53 2023 | Deriva de conceito | CD | 109,8K |
1. Preparação do ambiente 2. Treinamento TrafficLLM 2.1. Preparando o Ponto de Verificação Pré-treinado 2.2. Conjunto de dados de pré-processamento 2.3. Treinamento do Tokenizer de Domínio de Tráfego (Opcional) 2.4. Ajuste de instrução de linguagem neural 2.5. Ajuste de Tráfego Específico da Tarefa 2.6. Adaptação Extensível com PEFT (EA-PEFT) 3. Avaliando TrafficLLM 3.1. Preparando Pontos de Verificação e Dados 3.2. Avaliação em execuçãoÍndice:
1. Preparação do ambiente [Voltar ao início]
Clone o repositório e instale o ambiente necessário executando os seguintes comandos.
conda create -n tráfegollm python=3.9 conda activate tráfegollm# Clone nosso clone TrafficLLMgit https://github.com/ZGC-LLM-Safety/TrafficLLM.gitcd TrafficLLM# Instale as bibliotecas necessáriaspip install -r requisitos.txt# Se trainingpip instale rouge_chinese nltk jieba conjuntos de dados
O TrafficLLM emprega três técnicas principais: tokenização de domínio de tráfego para processar instruções e dados de tráfego, pipeline de ajuste de dois estágios para entender a semântica do texto e aprender padrões de tráfego em diferentes tarefas, e o EA-PEFT para atualizar parâmetros do modelo para adaptação a novos cenários.
O TrafficLLM é treinado com base em LLMs de código aberto existentes. Por favor, siga as instruções para preparar os pontos de verificação.
ChatGLM2 : Prepare o modelo básico ChatGLM, que é um LLM de código aberto com requisitos de implantação leves. Faça o download dos pesos aqui. Geralmente utilizamos o modelo v2 com parâmetros 6B.
Other LLMs : Para adaptar outros LLMs para tarefas de análise de tráfego, você pode reutilizar os dados de treinamento no repositório e modificar seus scripts de treinamento de acordo com as instruções oficiais. Por exemplo, Llama2 é necessário para registrar o novo conjunto de dados nas configurações.
Para extrair dados de treinamento adequados para o aprendizado LLM dos conjuntos de dados de tráfego brutos, projetamos extratores especializados para pré-processar conjuntos de dados de tráfego para diferentes tarefas. O código de pré-processamento contém os seguintes parâmetros para configuração.
input : o caminho do conjunto de dados de tráfego bruto (o caminho do diretório principal que contém subdiretórios rotulados. Cada subdiretório rotulado contém os arquivos .pcap brutos a serem pré-processados).
dataset_name : o nome do conjunto de dados de tráfego bruto (ajuda a avaliar se o nome foi registrado nos códigos do TrafficLLM).
traffic_task : tarefas de detecção ou tarefas de geração.
granularity : granularidade em nível de pacote ou nível de fluxo.
output_path : caminho do conjunto de dados de treinamento de saída.
output_name : nome do conjunto de dados de treinamento de saída.
Esta é uma instância para pré-processar conjuntos de dados de tráfego bruto para tarefas de detecção de tráfego em nível de pacote.
pré-processamento de CD python preprocess_dataset.py --input /Seu/Raw/Dataset/Caminho --dataset_name /Seu/Raw/Dataset/Nome --traffic_task detecção --pacote de granularidade --output_path /Sua/Saída/Dataset/Caminho --output_name /Seu /Output/Dataset/Name
TrafficLLM apresenta um tokenizador de domínio de tráfego para lidar com linguagem neural e dados de tráfego. Se você deseja treinar um tokenizer personalizado com seu próprio conjunto de dados, modifique model_name e data_path no código.
model_name : o caminho do modelo base que contém o tokenizer nativo.
data_path : os conjuntos de dados de treinamento extraídos do processo de pré-processamento.
Por favor, siga o comando para usar o código.
tokenização de CD python tráfego_tokenizer.py
Preparar dados: Os dados de ajuste de instrução de linguagem neural são nossos conjuntos de dados de instrução coletados para compreensão da tarefa de análise de tráfego.
Iniciar ajuste: Após as etapas mencionadas acima, você pode iniciar o ajuste do primeiro estágio usando tráfegollm_stage1.sh. Há um exemplo abaixo:
PRE_SEQ_LEN=128
LR=2e-2
NUM_GPUS=1exportar CUDA_VISIBLE_DEVICES=1
torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py
--do_train
--train_file ../datasets/instructions/instructions.json
--validation_file ../datasets/instructions/instructions.json
--preprocessing_num_workers 10
--prompt_column instrução
--resposta_coluna saída
--overwrite_cache
--cache_dir../cache
--model_name_or_path ../models/chatglm2/chatglm2-6b
--output_dir ../models/chatglm2/peft/instruction
--overwrite_output_dir
--max_source_length 1024
--max_target_length 32
--per_device_train_batch_size 1
--per_device_eval_batch_size 1
--gradient_accumulation_steps 16
--predict_with_generate
--max_steps 20000
--logging_steps 10
--save_steps 4000
--learning_rate $LR
--pre_seq_len $PRE_SEQ_LEN Preparar dados: Os conjuntos de dados de ajuste de tráfego específicos da tarefa são os conjuntos de dados de treinamento extraídos da etapa de pré-processamento para diferentes tarefas downstream.
Iniciar o ajuste: Após as etapas mencionadas acima, você poderá iniciar o ajuste do segundo estágio usando tráfegollm_stage2.sh. Há um exemplo abaixo:
PRE_SEQ_LEN=128
LR=2e-2
NUM_GPUS=1exportar CUDA_VISIBLE_DEVICES=1
torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS main.py
--do_train
--train_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_train.json
--validation_file ../datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_train.json
--preprocessing_num_workers 10
--prompt_column instrução
--resposta_coluna saída
--overwrite_cache
--cache_dir../cache
--model_name_or_path ../models/chatglm2/chatglm2-6b
--output_dir ../models/chatglm2/peft/ustc-tfc-2016-detection-packet
--overwrite_output_dir
--max_source_length 1024
--max_target_length 32
--per_device_train_batch_size 1
--per_device_eval_batch_size 1
--gradient_accumulation_steps 16
--predict_with_generate
--max_steps 20000
--logging_steps 10
--save_steps 4000
--learning_rate $LR
--pre_seq_len $PRE_SEQ_LEN TrafficLLM emprega EA-PEFT para organizar os modelos de ajuste fino de parâmetro efetivo (PEFT) com uma adaptação extensível, que pode ajudar TrafficLLM a se adaptar facilmente a novos ambientes. O adaptador TrafficLLM permite operações flexíveis para atualizar modelos antigos ou registrar novas tarefas.
model_name : o caminho do modelo base.
tuning_data : o novo conjunto de dados do ambiente.
adaptation_task : Atualizar ou registrar (atualizar modelos antigos ou registrar novas tarefas).
task_name : o nome da tarefa downstream a ser atualizada ou inserida.
Há um exemplo para atualizar o TrafficLLM com tarefas de detecção de tráfego de malware (MTD).
CD EA-PEFT python ea-peft.py --model_name /Your/Base/Model/Path --tuning_data /Your/New/Dataset/Path --adaptation_task update --task_name MTD
Pontos de verificação: você pode tentar avaliar o TrafficLLM usando seu próprio modelo ou nossos pontos de verificação lançados.
Dados: durante a etapa de pré-processamento, dividimos os conjuntos de dados de teste e construímos arquivos de rótulos para diferentes conjuntos de dados para avaliação. Consulte os códigos de pré-processamento.
Para medir a eficácia do TrafficLLM para diferentes tarefas downstream, execute os códigos de avaliação.
model_name : o caminho do modelo base.
traffic_task : tarefas de detecção ou tarefas de geração.
test_file : os conjuntos de dados de teste extraídos durante as etapas de pré-processamento.
label_file : o arquivo de rótulo extraído durante as etapas de pré-processamento.
ptuning_path : O caminho do modelo PEFT para avaliação específica de tarefa.
Há um exemplo para executar avaliação em tarefas MTD.
python avaliação.py --model_name /Sua/Base/Modelo/Caminho --traffic_task detecção --test_file datasets/ustc-tfc-2016/ustc-tfc-2016_detection_packet_test.json --label_file datasets/ustc-tfc-2016/ustc- tfc-2016_label.json --ptuning_path models/chatglm2/peft/ustc-tfc-2016-detection-packet/checkpoints-20000/
Você pode implantar o TrafficLLM em seu dispositivo local. Primeiro, configure o caminho do modelo em config.json para registrar os modelos PEFT coletados nas etapas de treinamento. Há um exemplo para registrar 6 tarefas no TrafficLLM:
{ "model_path": "models/chatglm2/chatglm2-6b/", "peft_path": "models/chatglm2/peft/", "peft_set": { "NLP": " instrução/ponto de verificação-8000/", "MTD": "ustc-tfc-2016-detection-packet/checkpoint-10000/", "BND": "iscx-botnet-2014-detection-packet/checkpoint-5000/", "WAD": "csic-2010-detection-packet /ponto de verificação-6000/", "AAD": "dapt-2020-detection-packet/checkpoint-20000/", "EVD": "iscx-vpn-2016-detection-packet/checkpoint-4000/", "TBD": "iscx-tor-2016-detection-packet /ponto de verificação-10000/"
}, "tarefas": { "Detecção de tráfego de malware": "MTD", "Detecção de botnet": "BND", "Detecção de ataque na Web": "WAD", "Detecção de ataque APT": "AAD", "Detecção de VPN criptografada ": "EVD", "Detecção de Comportamento Tor": "TBD"
}
} Então você deve adicionar o pré-prompt na função prepromt de inference.py e tráfegollm_server.py. O pré-prompt é o texto do prefixo usado nos dados de treinamento durante o ajuste de tráfego específico da tarefa.
Para conversar com o TrafficLLM no modo terminal, você pode executar o seguinte comando:
python inference.py --config=config.json --prompt="Seu texto de instrução +: + Dados de tráfego"
Você pode inicie a demonstração do site do TrafficLLM usando o seguinte comando:
streamlit run tráfegollm_server.py
Esta demonstração executa um servidor Web do TrafficLLM. Acesse http://Your-Server-IP:Port para conversar na caixa de chat.
Muito obrigado ao trabalho relacionado ChatGLM2 e Llama2 que serve de base para nossa estrutura e códigos. O design da construção do TrafficLLM é inspirado em ET-BERT e GraphGPT. Obrigado por seus maravilhosos trabalhos.
